Support for SELECT on nested JSON fields in Flat Collections

document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java

@@ -89,6 +89,7 @@
 import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
 import org.hypertrace.core.documentstore.expression.impl.FunctionExpression;
 import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
+import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
 import org.hypertrace.core.documentstore.expression.impl.KeyExpression;
 import org.hypertrace.core.documentstore.expression.impl.LogicalExpression;
 import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
@@ -3848,6 +3849,138 @@ void testFlatPostgresCollectionBooleanArrayFilter(String dataStoreName) throws I
       assertDocsAndSizeEqualWithoutOrder(
           dataStoreName, resultIterator, "query/flat_boolean_array_filter_response.json", 2);
     }
+
+    /**
+     * Tests selection of JSONB nested fields using JsonIdentifierExpression on flat collection.
+     * Validates selecting simple nested fields, deeply nested fields, and JSONB arrays without any
+     * filters.
+     */
+    @ParameterizedTest
+    @ArgumentsSource(PostgresProvider.class)
+    void testFlatCollectionNestedJsonSelections(String dataStoreName) throws IOException {
+      Datastore datastore = datastoreMap.get(dataStoreName);
+      Collection flatCollection =
+          datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);
+
+      // Test 1: Select nested STRING field from JSONB column (props.brand)
+      Query brandSelectionQuery =
+          Query.builder().addSelection(JsonIdentifierExpression.of("props", "brand")).build();
+
+      Iterator<Document> brandIterator = flatCollection.find(brandSelectionQuery);
+      assertDocsAndSizeEqualWithoutOrder(
+          dataStoreName, brandIterator, "query/flat_jsonb_brand_selection_response.json", 10);
+
+      // Test 2: Select deeply nested STRING field from JSONB column (props.seller.address.city)
+      Query citySelectionQuery =
+          Query.builder()
+              .addSelection(JsonIdentifierExpression.of("props", "seller", "address", "city"))
+              .build();
+
+      Iterator<Document> cityIterator = flatCollection.find(citySelectionQuery);
+      assertDocsAndSizeEqualWithoutOrder(
+          dataStoreName, cityIterator, "query/flat_jsonb_city_selection_response.json", 10);
+
+      // Test 3: Select STRING_ARRAY field from JSONB column (props.colors)
+      Query colorsSelectionQuery =
+          Query.builder().addSelection(JsonIdentifierExpression.of("props", "colors")).build();
+
+      Iterator<Document> colorsIterator = flatCollection.find(colorsSelectionQuery);
+      assertDocsAndSizeEqualWithoutOrder(
+          dataStoreName, colorsIterator, "query/flat_jsonb_colors_selection_response.json", 10);
+    }
+
+    @ParameterizedTest
+    @ArgumentsSource(PostgresProvider.class)
+    void testFlatVsNestedCollectionNestedFieldSelections(String dataStoreName) throws IOException {
+      Datastore datastore = datastoreMap.get(dataStoreName);
+
+      Collection nestedCollection = datastore.getCollection(COLLECTION_NAME);
+      Collection flatCollection =
+          datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);
+
+      // Test 1: Select nested field - props.brand
+      // Nested collection uses dot notation
+      Query nestedBrandQuery =
+          Query.builder()
+              .addSelection(IdentifierExpression.of("item"))
+              .addSelection(IdentifierExpression.of("props.brand"), "brand")
+              .addSort(IdentifierExpression.of("item"), ASC)
+              .build();
+
+      // Flat collection uses JsonIdentifierExpression
+      // Filter to exclude docs 9-10 to match nested collection dataset
+      Query flatBrandQuery =
+          Query.builder()
+              .addSelection(IdentifierExpression.of("item"))
+              .addSelection(JsonIdentifierExpression.of("props", "brand"), "brand")
+              .addSort(IdentifierExpression.of("item"), ASC)
+              .setFilter(
+                  RelationalExpression.of(
+                      IdentifierExpression.of("_id"), LTE, ConstantExpression.of(8)))
+              .build();
+
+      // Assert both match the expected response
+      Iterator<Document> nestedBrandIterator = nestedCollection.find(nestedBrandQuery);
+      assertDocsAndSizeEqual(dataStoreName, nestedBrandIterator, "query/brand_response.json", 8);
+
+      Iterator<Document> flatBrandIterator = flatCollection.find(flatBrandQuery);
+      assertDocsAndSizeEqual(dataStoreName, flatBrandIterator, "query/brand_response.json", 8);
+
+      // Test 2: Select nested JSON array - props.colors
+      Query nestedColorsQuery =
+          Query.builder()
+              .addSelection(IdentifierExpression.of("item"))
+              .addSelection(IdentifierExpression.of("props.colors"), "colors")
+              .addSort(IdentifierExpression.of("item"), ASC)
+              .build();
+
+      // Filter to exclude docs 9-10 to match nested collection dataset
+      Query flatColorsQuery =
+          Query.builder()
+              .addSelection(IdentifierExpression.of("item"))
+              .addSelection(JsonIdentifierExpression.of("props", "colors"), "colors")
+              .addSort(IdentifierExpression.of("item"), ASC)
+              .setFilter(
+                  RelationalExpression.of(
+                      IdentifierExpression.of("_id"), LTE, ConstantExpression.of(8)))
+              .build();
+
+      // Assert both match the expected response
+      Iterator<Document> nestedColorsIterator = nestedCollection.find(nestedColorsQuery);
+      assertDocsAndSizeEqual(dataStoreName, nestedColorsIterator, "query/colors_response.json", 8);
+
+      Iterator<Document> flatColorsIterator = flatCollection.find(flatColorsQuery);
+      assertDocsAndSizeEqual(dataStoreName, flatColorsIterator, "query/colors_response.json", 8);
+
+      // Test 3: Select nested field WITHOUT alias - should preserve full nested structure
+      Query nestedBrandNoAliasQuery =
+          Query.builder()
+              .addSelection(IdentifierExpression.of("item"))
+              .addSelection(IdentifierExpression.of("props.brand"))
+              .addSort(IdentifierExpression.of("item"), ASC)
+              .build();
+
+      // Filter to exclude docs 9-10 to match nested collection dataset
+      Query flatBrandNoAliasQuery =
+          Query.builder()
+              .addSelection(IdentifierExpression.of("item"))
+              .addSelection(JsonIdentifierExpression.of("props", "brand"))
+              .addSort(IdentifierExpression.of("item"), ASC)
+              .setFilter(
+                  RelationalExpression.of(
+                      IdentifierExpression.of("_id"), LTE, ConstantExpression.of(8)))
+              .build();
+
+      // Assert both match the expected response with nested structure
+      Iterator<Document> nestedBrandNoAliasIterator =
+          nestedCollection.find(nestedBrandNoAliasQuery);
+      assertDocsAndSizeEqual(
+          dataStoreName, nestedBrandNoAliasIterator, "query/no_alias_response.json", 8);
+
+      Iterator<Document> flatBrandNoAliasIterator = flatCollection.find(flatBrandNoAliasQuery);
+      assertDocsAndSizeEqual(
+          dataStoreName, flatBrandNoAliasIterator, "query/no_alias_response.json", 8);
+    }
   }
 
   @Nested

document-store/src/integrationTest/resources/query/brand_response.json

@@ -0,0 +1,29 @@
+[
+  {
+    "item": "Comb"
+  },
+  {
+    "item": "Comb"
+  },
+  {
+    "item": "Mirror"
+  },
+  {
+    "item": "Shampoo",
+    "brand": "Sunsilk"
+  },
+  {
+    "item": "Shampoo"
+  },
+  {
+    "item": "Soap"
+  },
+  {
+    "item": "Soap",
+    "brand": "Lifebuoy"
+  },
+  {
+    "item": "Soap",
+    "brand": "Dettol"
+  }
+]

document-store/src/integrationTest/resources/query/colors_response.json

@@ -0,0 +1,39 @@
+[
+  {
+    "item": "Comb"
+  },
+  {
+    "colors": [
+    ],
+    "item": "Comb"
+  },
+  {
+    "item": "Mirror"
+  },
+  {
+    "colors": [
+      "Black"
+    ],
+    "item": "Shampoo"
+  },
+  {
+    "item": "Shampoo"
+  },
+  {
+    "item": "Soap"
+  },
+  {
+    "colors": [
+      "Orange",
+      "Blue"
+    ],
+    "item": "Soap"
+  },
+  {
+    "colors": [
+      "Blue",
+      "Green"
+    ],
+    "item": "Soap"
+  }
+]

document-store/src/integrationTest/resources/query/flat_brand_response.json

@@ -0,0 +1,35 @@
+[
+  {
+    "item": "Bottle"
+  },
+  {
+    "item": "Comb"
+  },
+  {
+    "item": "Comb"
+  },
+  {
+    "item": "Cup"
+  },
+  {
+    "item": "Mirror"
+  },
+  {
+    "item": "Shampoo"
+  },
+  {
+    "item": "Shampoo",
+    "brand": "Sunsilk"
+  },
+  {
+    "item": "Soap",
+    "brand": "Lifebuoy"
+  },
+  {
+    "item": "Soap"
+  },
+  {
+    "item": "Soap",
+    "brand": "Dettol"
+  }
+]

document-store/src/integrationTest/resources/query/flat_colors_response.json

@@ -0,0 +1,44 @@
+[
+  {
+    "item": "Bottle"
+  },
+  {
+    "item": "Comb"
+  },
+  {
+    "colors": [],
+    "item": "Comb"
+  },
+  {
+    "item": "Cup"
+  },
+  {
+    "item": "Mirror"
+  },
+  {
+    "item": "Shampoo"
+  },
+  {
+    "colors": [
+      "Black"
+    ],
+    "item": "Shampoo"
+  },
+  {
+    "colors": [
+      "Orange",
+      "Blue"
+    ],
+    "item": "Soap"
+  },
+  {
+    "item": "Soap"
+  },
+  {
+    "colors": [
+      "Blue",
+      "Green"
+    ],
+    "item": "Soap"
+  }
+]

document-store/src/integrationTest/resources/query/flat_jsonb_add_response.json

@@ -0,0 +1,10 @@
+[
+  {"item": "Soap", "rating_plus_discount": 14.75},
+  {"item": "Mirror"},
+  {"item": "Shampoo", "rating_plus_discount": 9.3},
+  {"item": "Shampoo"},
+  {"item": "Soap", "rating_plus_discount": -4.55},
+  {"item": "Comb"},
+  {"item": "Comb", "rating_plus_discount": 3.5},
+  {"item": "Soap"}
+]

document-store/src/integrationTest/resources/query/flat_jsonb_brand_selection_response.json

@@ -0,0 +1,24 @@
+[
+  {
+    "props": {
+      "brand": "Dettol"
+    }
+  },
+  {},
+  {
+    "props": {
+      "brand": "Sunsilk"
+    }
+  },
+  {},
+  {
+    "props": {
+      "brand": "Lifebuoy"
+    }
+  },
+  {},
+  {},
+  {},
+  {},
+  {}
+]

document-store/src/integrationTest/resources/query/flat_jsonb_city_selection_response.json

@@ -0,0 +1,47 @@
+[
+  {
+    "props": {
+      "seller": {
+        "address": {
+          "city": "Mumbai"
+        }
+      }
+    }
+  },
+  {
+  },
+  {
+    "props": {
+      "seller": {
+        "address": {
+          "city": "Mumbai"
+        }
+      }
+    }
+  },
+  {
+  },
+  {
+    "props": {
+      "seller": {
+        "address": {
+          "city": "Kolkata"
+        }
+      }
+    }
+  },
+  {
+  },
+  {
+    "props": {
+      "seller": {
+        "address": {
+          "city": "Kolkata"
+        }
+      }
+    }
+  },
+  {},
+  {},
+  {}
+]