Adds support for index Document properties in AppSearch.Document.

Bug: 143789408
Test: atest FrameworksCoreTests:android.app.appsearch
Change-Id: Ie5d2b7a3c929cc5d0a6edee0a0a01eb06e9a4e3a

3 files changed, 88 insertions, 2 deletions

diff --git a/apex/appsearch/framework/java/android/app/appsearch/AppSearchDocument.java b/apex/appsearch/framework/java/android/app/appsearch/AppSearchDocument.java
index 66a90c25f7c4..ff0f0dda55b9 100644
--- a/apex/appsearch/framework/java/android/app/appsearch/AppSearchDocument.java
+++ b/apex/appsearch/framework/java/android/app/appsearch/AppSearchDocument.java
@@ -301,6 +301,23 @@ public class AppSearchDocument {
         return propertyArray[0];
     }
 
+    /**
+     * Retrieve a {@link AppSearchDocument} value by key.
+     *
+     * @param key The key to look for.
+     * @return The first {@link AppSearchDocument} associated with the given key or {@code null} if
+     *         there is no such key or the value is of a different type.
+     */
+    @Nullable
+    public AppSearchDocument getPropertyDocument(@NonNull String key) {
+        AppSearchDocument[] propertyArray = getPropertyDocumentArray(key);
+        if (ArrayUtils.isEmpty(propertyArray)) {
+            return null;
+        }
+        warnIfSinglePropertyTooLong("Document", key, propertyArray.length);
+        return propertyArray[0];
+    }
+
     /** Prints a warning to logcat if the given propertyLength is greater than 1. */
     private static void warnIfSinglePropertyTooLong(
             @NonNull String propertyType, @NonNull String key, int propertyLength) {
@@ -377,6 +394,18 @@ public class AppSearchDocument {
     }
 
     /**
+     * Retrieve a repeated {@link AppSearchDocument} property by key.
+     *
+     * @param key The key to look for.
+     * @return The {@link AppSearchDocument[]} associated with the given key, or {@code null} if no
+     *         value is set or the value is of a different type.
+     */
+    @Nullable
+    public AppSearchDocument[] getPropertyDocumentArray(@NonNull String key) {
+        return getAndCastPropertyArray(key, AppSearchDocument[].class);
+    }
+
+    /**
      * Gets a repeated property of the given key, and casts it to the given class type, which
      * must be an array class type.
      */
@@ -451,7 +480,7 @@ public class AppSearchDocument {
         }
 
         /**
-         * Set the score of the {@link AppSearchDocument}.
+         * Sets the score of the {@link AppSearchDocument}.
          *
          * <p>The score is a query-independent measure of the document's quality, relative to
          * other {@link AppSearchDocument}s of the same type.
@@ -563,6 +592,19 @@ public class AppSearchDocument {
             return mBuilderTypeInstance;
         }
 
+        /**
+         * Sets one or multiple {@link AppSearchDocument} values for a property, replacing its
+         * previous values.
+         *
+         * @param key The key associated with the {@code values}.
+         * @param values The {@link AppSearchDocument} values of the property.
+         */
+        @NonNull
+        public BuilderType setProperty(@NonNull String key, @NonNull AppSearchDocument... values) {
+            putInPropertyMap(key, values);
+            return mBuilderTypeInstance;
+        }
+
         private void putInPropertyMap(@NonNull String key, @NonNull String[] values)
                 throws IllegalArgumentException {
             Objects.requireNonNull(key);
@@ -608,6 +650,17 @@ public class AppSearchDocument {
             mProperties.put(key, values);
         }
 
+        private void putInPropertyMap(@NonNull String key, @NonNull AppSearchDocument[] values) {
+            Objects.requireNonNull(key);
+            Objects.requireNonNull(values);
+            for (int i = 0; i < values.length; i++) {
+                if (values[i] == null) {
+                    throw new IllegalArgumentException("The document at " + i + " is null.");
+                }
+            }
+            validateRepeatedPropertyLength(key, values.length);
+            mProperties.put(key, values);
+        }
 
         private static void validateRepeatedPropertyLength(@NonNull String key, int length) {
             if (length == 0) {
@@ -650,6 +703,10 @@ public class AppSearchDocument {
                     for (String value : (String[]) values) {
                         propertyProto.addStringValues(value);
                     }
+                } else if (values instanceof AppSearchDocument[]) {
+                    for (AppSearchDocument value : (AppSearchDocument[]) values) {
+                        propertyProto.addDocumentValues(value.getProto());
+                    }
                 } else if (values instanceof byte[][]) {
                     for (byte[] value : (byte[][]) values) {
                         propertyProto.addBytesValues(ByteString.copyFrom(value));
diff --git a/core/tests/coretests/src/android/app/appsearch/AppSearchDocumentTest.java b/core/tests/coretests/src/android/app/appsearch/AppSearchDocumentTest.java
index f2911717c915..4a4f13676cb3 100644
--- a/core/tests/coretests/src/android/app/appsearch/AppSearchDocumentTest.java
+++ b/core/tests/coretests/src/android/app/appsearch/AppSearchDocumentTest.java
@@ -37,6 +37,12 @@ import java.util.List;
 public class AppSearchDocumentTest {
     private static final byte[] sByteArray1 = new byte[]{(byte) 1, (byte) 2, (byte) 3};
     private static final byte[] sByteArray2 = new byte[]{(byte) 4, (byte) 5, (byte) 6};
+    private static final AppSearchDocument sDocumentProperties1 = new AppSearchDocument
+            .Builder("sDocumentProperties1", "sDocumentPropertiesSchemaType1")
+            .build();
+    private static final AppSearchDocument sDocumentProperties2 = new AppSearchDocument
+            .Builder("sDocumentProperties2", "sDocumentPropertiesSchemaType2")
+            .build();
 
     @Test
     public void testDocumentEquals_Identical() {
@@ -48,6 +54,7 @@ public class AppSearchDocumentTest {
                 .setProperty("booleanKey1", true, false, true)
                 .setProperty("stringKey1", "test-value1", "test-value2", "test-value3")
                 .setProperty("byteKey1", sByteArray1, sByteArray2)
+                .setProperty("documentKey1", sDocumentProperties1, sDocumentProperties2)
                 .build();
         AppSearchDocument document2 = new AppSearchDocument.Builder("uri1", "schemaType1")
                 .setCreationTimestampMillis(5L)
@@ -57,6 +64,7 @@ public class AppSearchDocumentTest {
                 .setProperty("booleanKey1", true, false, true)
                 .setProperty("stringKey1", "test-value1", "test-value2", "test-value3")
                 .setProperty("byteKey1", sByteArray1, sByteArray2)
+                .setProperty("documentKey1", sDocumentProperties1, sDocumentProperties2)
                 .build();
         assertThat(document1).isEqualTo(document2);
         assertThat(document1.hashCode()).isEqualTo(document2.hashCode());
@@ -70,6 +78,7 @@ public class AppSearchDocumentTest {
                 .setProperty("byteKey1", sByteArray1, sByteArray2)
                 .setProperty("doubleKey1", 1.0, 2.0, 3.0)
                 .setProperty("booleanKey1", true, false, true)
+                .setProperty("documentKey1", sDocumentProperties1, sDocumentProperties2)
                 .setProperty("stringKey1", "test-value1", "test-value2", "test-value3")
                 .build();
 
@@ -77,6 +86,7 @@ public class AppSearchDocumentTest {
         AppSearchDocument document2 = new AppSearchDocument.Builder("uri1", "schemaType1")
                 .setCreationTimestampMillis(5L)
                 .setProperty("booleanKey1", true, false, true)
+                .setProperty("documentKey1", sDocumentProperties1, sDocumentProperties2)
                 .setProperty("stringKey1", "test-value1", "test-value2", "test-value3")
                 .setProperty("doubleKey1", 1.0, 2.0, 3.0)
                 .setProperty("byteKey1", sByteArray1, sByteArray2)
@@ -128,7 +138,9 @@ public class AppSearchDocumentTest {
                 .setProperty("doubleKey1", 1.0)
                 .setProperty("booleanKey1", true)
                 .setProperty("stringKey1", "test-value1")
-                .setProperty("byteKey1", sByteArray1).build();
+                .setProperty("byteKey1", sByteArray1)
+                .setProperty("documentKey1", sDocumentProperties1)
+                .build();
         assertThat(document.getUri()).isEqualTo("uri1");
         assertThat(document.getTtlMillis()).isEqualTo(1L);
         assertThat(document.getSchemaType()).isEqualTo("schemaType1");
@@ -140,6 +152,7 @@ public class AppSearchDocumentTest {
         assertThat(document.getPropertyString("stringKey1")).isEqualTo("test-value1");
         assertThat(document.getPropertyBytes("byteKey1"))
                 .asList().containsExactly((byte) 1, (byte) 2, (byte) 3);
+        assertThat(document.getPropertyDocument("documentKey1")).isEqualTo(sDocumentProperties1);
     }
 
     @Test
@@ -151,6 +164,7 @@ public class AppSearchDocumentTest {
                 .setProperty("booleanKey1", true, false, true)
                 .setProperty("stringKey1", "test-value1", "test-value2", "test-value3")
                 .setProperty("byteKey1", sByteArray1, sByteArray2)
+                .setProperty("documentKey1", sDocumentProperties1, sDocumentProperties2)
                 .build();
 
         assertThat(document.getUri()).isEqualTo("uri1");
@@ -164,6 +178,8 @@ public class AppSearchDocumentTest {
                 .containsExactly("test-value1", "test-value2", "test-value3");
         assertThat(document.getPropertyBytesArray("byteKey1")).asList()
                 .containsExactly(sByteArray1, sByteArray2);
+        assertThat(document.getPropertyDocumentArray("documentKey1")).asList()
+                .containsExactly(sDocumentProperties1, sDocumentProperties2);
     }
 
     @Test
@@ -211,6 +227,7 @@ public class AppSearchDocumentTest {
                 .setProperty("booleanKey1", true)
                 .setProperty("stringKey1", "test-value1")
                 .setProperty("byteKey1", sByteArray1)
+                .setProperty("documentKey1", sDocumentProperties1)
                 .build();
 
         // Create the Document proto. Need to sort the property order by key.
@@ -232,6 +249,9 @@ public class AppSearchDocumentTest {
         propertyProtoMap.put("byteKey1",
                 PropertyProto.newBuilder().setName("byteKey1").addBytesValues(
                         ByteString.copyFrom(sByteArray1)));
+        propertyProtoMap.put("documentKey1",
+                PropertyProto.newBuilder().setName("documentKey1")
+                        .addDocumentValues(sDocumentProperties1.getProto()));
         List<String> sortedKey = new ArrayList<>(propertyProtoMap.keySet());
         Collections.sort(sortedKey);
         for (String key : sortedKey) {
diff --git a/core/tests/coretests/src/android/app/appsearch/impl/CustomerDocumentTest.java b/core/tests/coretests/src/android/app/appsearch/impl/CustomerDocumentTest.java
index 24d8b44ace5f..b29483c2e3b3 100644
--- a/core/tests/coretests/src/android/app/appsearch/impl/CustomerDocumentTest.java
+++ b/core/tests/coretests/src/android/app/appsearch/impl/CustomerDocumentTest.java
@@ -37,6 +37,12 @@ public class CustomerDocumentTest {
 
     private static byte[] sByteArray1 = new byte[]{(byte) 1, (byte) 2, (byte) 3};
     private static byte[] sByteArray2 = new byte[]{(byte) 4, (byte) 5, (byte) 6};
+    private static AppSearchDocument sDocumentProperties1 = new AppSearchDocument
+            .Builder("sDocumentProperties1", "sDocumentPropertiesSchemaType1")
+            .build();
+    private static AppSearchDocument sDocumentProperties2 = new AppSearchDocument
+            .Builder("sDocumentProperties2", "sDocumentPropertiesSchemaType2")
+            .build();
 
     @Test
     public void testBuildCustomerDocument() {
@@ -48,6 +54,7 @@ public class CustomerDocumentTest {
                 .setProperty("booleanKey1", true, false, true)
                 .setProperty("stringKey1", "test-value1", "test-value2", "test-value3")
                 .setProperty("byteKey1", sByteArray1, sByteArray2)
+                .setProperty("documentKey1", sDocumentProperties1, sDocumentProperties2)
                 .build();
 
         assertThat(customerDocument.getUri()).isEqualTo("uri1");
@@ -64,6 +71,8 @@ public class CustomerDocumentTest {
                 .containsExactly("test-value1", "test-value2", "test-value3");
         assertThat(customerDocument.getPropertyBytesArray("byteKey1")).asList()
                 .containsExactly(sByteArray1, sByteArray2);
+        assertThat(customerDocument.getPropertyDocumentArray("documentKey1")).asList()
+                .containsExactly(sDocumentProperties1, sDocumentProperties2);
     }
 
     /**