From 58cae6f69551898f5d8d277a765b39701b4e6ecb Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Fri, 19 Jun 2026 15:48:02 +0530 Subject: [PATCH 01/12] Changes to enhance existing length op --- .../documentstore/DocStoreQueryV1Test.java | 49 +++++++++++++++++++ .../parser/MongoFunctionExpressionParser.java | 5 ++ .../PostgresFunctionExpressionVisitor.java | 5 +- .../query/v1/PostgresQueryParserTest.java | 6 +-- .../mongo/pipeline/distinct_count.json | 7 ++- .../resources/mongo/pipeline/field_count.json | 7 ++- ...imple_sort_with_aggregation_selection.json | 7 ++- .../test/resources/mongo/pipeline/simple.json | 7 ++- .../mongo/pipeline/with_projections.json | 7 ++- 9 files changed, 91 insertions(+), 9 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java index 5c9f8930f..b4d677b35 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java @@ -919,6 +919,55 @@ public void testAggregateWithMultipleGroupingLevels(String dataStoreName) throws testCountApi(dataStoreName, query, "query/multi_level_grouping_response.json"); } + @ParameterizedTest + @ArgumentsSource(MongoProvider.class) + public void testSortByListSizeWithMissingField(String dataStoreName) throws IOException { + Datastore datastore = datastoreMap.get(dataStoreName); + String collectionName = "list_size_sort_collection"; + datastore.deleteCollection(collectionName); + datastore.createCollection(collectionName, null); + Collection collection = datastore.getCollection(collectionName); + + collection.upsert( + new SingleValueKey(TENANT_ID, "three"), + new JSONDocument("{\"item\":\"three\",\"tags\":[\"a\",\"b\",\"c\"]}")); + collection.upsert( + new SingleValueKey(TENANT_ID, "one"), + new JSONDocument("{\"item\":\"one\",\"tags\":[\"x\"]}")); + // Document intentionally missing the "tags" field; LENGTH must resolve to 0 instead of failing + collection.upsert( + new SingleValueKey(TENANT_ID, "none"), new JSONDocument("{\"item\":\"none\"}")); + + Query query = + Query.builder() + .addSelection(IdentifierExpression.of("item")) + .addSelection( + FunctionExpression.builder() + .operator(LENGTH) + .operand(IdentifierExpression.of("tags")) + .build(), + "tag_count") + .addSort(IdentifierExpression.of("tag_count"), ASC) + .build(); + + Iterator resultDocs = collection.aggregate(query); + List> results = new ArrayList<>(); + while (resultDocs.hasNext()) { + results.add(Utils.convertDocumentToMap(resultDocs.next())); + } + + assertEquals(3, results.size()); + // Document without the "tags" field counts as 0 and sorts first in ascending order + assertEquals("none", results.get(0).get("item")); + assertEquals(0, ((Number) results.get(0).get("tag_count")).intValue()); + assertEquals("one", results.get(1).get("item")); + assertEquals(1, ((Number) results.get(1).get("tag_count")).intValue()); + assertEquals("three", results.get(2).get("item")); + assertEquals(3, ((Number) results.get(2).get("tag_count")).intValue()); + + datastore.deleteCollection(collectionName); + } + @ParameterizedTest @ArgumentsSource(AllProvider.class) public void testAggregateWithFunctionalLeftHandSideFilter(final String dataStoreName) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionExpressionParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionExpressionParser.java index a9c27f96f..ddac4aa10 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionExpressionParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionExpressionParser.java @@ -69,6 +69,11 @@ Map parse(final FunctionExpression expression) { if (numArgs == 1) { Object value = expression.getOperands().get(0).accept(parser); + // $size fails when the operand resolves to a missing/absent (or null) field. Default such a + // value to an empty array so that LENGTH of an absent field is 0 instead of throwing an error. + if (operator == LENGTH) { + value = Map.of("$ifNull", List.of(value, List.of())); + } return Map.of(key, value); } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java index f1b0f2ee4..af267f57b 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java @@ -51,7 +51,10 @@ public String visit(final FunctionExpression expression) { if (numArgs == 1) { String parsedExpression = getParsedExpression(expression.getOperands().get(0)); return expression.getOperator().equals(FunctionOperator.LENGTH) - ? String.format("ARRAY_LENGTH( %s, %s )", parsedExpression, ARRAY_DIMENSION) + // COALESCE so that LENGTH of an absent/empty array resolves to 0 instead of NULL, keeping + // parity with the Mongo backend ($ifNull) behaviour. + ? String.format( + "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", parsedExpression, ARRAY_DIMENSION) : String.format("%s( %s )", expression.getOperator(), parsedExpression); } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java index 5a42e1cd7..c78337aee 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParserTest.java @@ -395,7 +395,7 @@ void testAggregationExpressionDistinctCount() { assertEquals( "SELECT ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)) AS \"qty_distinct\", " - + "ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ) AS \"qty_distinct_length\" " + + "COALESCE( ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ), 0 ) AS \"qty_distinct_length\" " + "FROM \"testCollection\" " + "WHERE CAST (document->>'price' AS NUMERIC) = ? " + "GROUP BY document->'item'", @@ -435,10 +435,10 @@ void testAggregateWithMultipleGroupingLevels() { assertEquals( "SELECT document->'item' AS \"item\", document->'price' AS \"price\", " + "ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)) AS \"quantities\", " - + "ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ) AS \"num_quantities\" " + + "COALESCE( ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ), 0 ) AS \"num_quantities\" " + "FROM \"testCollection\" " + "GROUP BY document->'item',document->'price' " - + "HAVING ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ) = ? " + + "HAVING COALESCE( ARRAY_LENGTH( ARRAY_AGG(DISTINCT CAST (document->>'quantity' AS NUMERIC)), 1 ), 0 ) = ? " + "ORDER BY document->'item' DESC NULLS LAST", sql); diff --git a/document-store/src/test/resources/mongo/pipeline/distinct_count.json b/document-store/src/test/resources/mongo/pipeline/distinct_count.json index 27c256b93..bdc99d184 100644 --- a/document-store/src/test/resources/mongo/pipeline/distinct_count.json +++ b/document-store/src/test/resources/mongo/pipeline/distinct_count.json @@ -19,7 +19,12 @@ { "$project": { "section_count": { - "$size": "$section_count" + "$size": { + "$ifNull": [ + "$section_count", + [] + ] + } } } } diff --git a/document-store/src/test/resources/mongo/pipeline/field_count.json b/document-store/src/test/resources/mongo/pipeline/field_count.json index bd7def235..3c8c9c5f9 100644 --- a/document-store/src/test/resources/mongo/pipeline/field_count.json +++ b/document-store/src/test/resources/mongo/pipeline/field_count.json @@ -10,7 +10,12 @@ { "$project": { "total": { - "$size": "$total" + "$size": { + "$ifNull": [ + "$total", + [] + ] + } } } } diff --git a/document-store/src/test/resources/mongo/pipeline/optimize_sorts_simple_sort_with_aggregation_selection.json b/document-store/src/test/resources/mongo/pipeline/optimize_sorts_simple_sort_with_aggregation_selection.json index ce2253af5..ee42beb4f 100644 --- a/document-store/src/test/resources/mongo/pipeline/optimize_sorts_simple_sort_with_aggregation_selection.json +++ b/document-store/src/test/resources/mongo/pipeline/optimize_sorts_simple_sort_with_aggregation_selection.json @@ -10,7 +10,12 @@ { "$project": { "total": { - "$size": "$total" + "$size": { + "$ifNull": [ + "$total", + [] + ] + } } } }, diff --git a/document-store/src/test/resources/mongo/pipeline/simple.json b/document-store/src/test/resources/mongo/pipeline/simple.json index f07ec7a03..97394ffe8 100644 --- a/document-store/src/test/resources/mongo/pipeline/simple.json +++ b/document-store/src/test/resources/mongo/pipeline/simple.json @@ -10,7 +10,12 @@ { "$project": { "total": { - "$size": "$total" + "$size": { + "$ifNull": [ + "$total", + [] + ] + } } } } diff --git a/document-store/src/test/resources/mongo/pipeline/with_projections.json b/document-store/src/test/resources/mongo/pipeline/with_projections.json index cbc11066d..dec1105d9 100644 --- a/document-store/src/test/resources/mongo/pipeline/with_projections.json +++ b/document-store/src/test/resources/mongo/pipeline/with_projections.json @@ -11,7 +11,12 @@ "$project": { "name": 1, "total": { - "$size": "$total" + "$size": { + "$ifNull": [ + "$total", + [] + ] + } } } } From 317cc257c2b91febe14fe6f325faa012cb2c5fe8 Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Fri, 19 Jun 2026 16:33:57 +0530 Subject: [PATCH 02/12] nit --- .../org/hypertrace/core/documentstore/DocStoreQueryV1Test.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java index b4d677b35..a1346bc08 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java @@ -920,7 +920,7 @@ public void testAggregateWithMultipleGroupingLevels(String dataStoreName) throws } @ParameterizedTest - @ArgumentsSource(MongoProvider.class) + @ArgumentsSource(AllProvider.class) public void testSortByListSizeWithMissingField(String dataStoreName) throws IOException { Datastore datastore = datastoreMap.get(dataStoreName); String collectionName = "list_size_sort_collection"; From 33538738c04dbac588f04d6ce6042d5663096bbf Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Fri, 19 Jun 2026 16:41:33 +0530 Subject: [PATCH 03/12] nit --- .../mongo/query/parser/MongoFunctionExpressionParser.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionExpressionParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionExpressionParser.java index ddac4aa10..99e587a84 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionExpressionParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoFunctionExpressionParser.java @@ -70,7 +70,8 @@ Map parse(final FunctionExpression expression) { if (numArgs == 1) { Object value = expression.getOperands().get(0).accept(parser); // $size fails when the operand resolves to a missing/absent (or null) field. Default such a - // value to an empty array so that LENGTH of an absent field is 0 instead of throwing an error. + // value to an empty array so that LENGTH of an absent field is 0 instead of throwing an + // error. if (operator == LENGTH) { value = Map.of("$ifNull", List.of(value, List.of())); } From 6abda6d0e90eefd428b3bc9d5f81f6334f93ab38 Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Fri, 19 Jun 2026 22:51:07 +0530 Subject: [PATCH 04/12] fix: use jsonb_array_length for raw jsonb fields in LENGTH operator When sorting on an array field that is missing or empty, the Postgres backend now correctly returns 0. Uses jsonb_array_length for raw jsonb field access and ARRAY_LENGTH for native PG arrays from aggregations. Co-Authored-By: Claude Opus 4.6 --- .../PostgresFunctionExpressionVisitor.java | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java index af267f57b..78bf80a95 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java @@ -49,13 +49,11 @@ public String visit(final FunctionExpression expression) { } if (numArgs == 1) { + if (expression.getOperator().equals(FunctionOperator.LENGTH)) { + return buildLengthExpression(expression.getOperands().get(0)); + } String parsedExpression = getParsedExpression(expression.getOperands().get(0)); - return expression.getOperator().equals(FunctionOperator.LENGTH) - // COALESCE so that LENGTH of an absent/empty array resolves to 0 instead of NULL, keeping - // parity with the Mongo backend ($ifNull) behaviour. - ? String.format( - "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", parsedExpression, ARRAY_DIMENSION) - : String.format("%s( %s )", expression.getOperator(), parsedExpression); + return String.format("%s( %s )", expression.getOperator(), parsedExpression); } Collector collector = @@ -86,6 +84,20 @@ private Collector getCollectorForFunctionOperator(FunctionOperator operator) { String.format("Query operation:%s not supported", operator)); } + private String buildLengthExpression(final SelectTypeExpression operand) { + Optional identifier = Optional.ofNullable(operand.accept(identifierExpressionVisitor)); + Optional resolvedSelection = + identifier.map(v -> getPostgresQueryParser().getPgSelections().get(v)); + if (resolvedSelection.isPresent()) { + // Operand resolved to a prior selection (e.g. ARRAY_AGG) which produces a native PG array. + return String.format( + "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", resolvedSelection.get(), ARRAY_DIMENSION); + } + // Raw jsonb field access — use jsonb_array_length which operates on jsonb arrays directly. + String parsedExpression = operand.accept(selectTypeExpressionVisitor); + return String.format("COALESCE( jsonb_array_length( %s ), 0 )", parsedExpression); + } + private String getParsedExpression(final SelectTypeExpression expression) { Optional identifier = Optional.ofNullable(expression.accept(identifierExpressionVisitor)); From ebdb64c741c591f8ebe359f663598119bae7838f Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Fri, 19 Jun 2026 22:57:37 +0530 Subject: [PATCH 05/12] fix: use jsonb field accessor without cast for jsonb_array_length The previous approach used PostgresDataAccessorIdentifierExpressionVisitor which casts jsonb to numeric via ->>. jsonb_array_length requires a raw jsonb value (via -> accessor), so switch to PostgresFieldIdentifierExpressionVisitor. Co-Authored-By: Claude Opus 4.6 --- .../query/v1/vistors/PostgresFunctionExpressionVisitor.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java index 78bf80a95..9cb85e29c 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java @@ -89,12 +89,12 @@ private String buildLengthExpression(final SelectTypeExpression operand) { Optional resolvedSelection = identifier.map(v -> getPostgresQueryParser().getPgSelections().get(v)); if (resolvedSelection.isPresent()) { - // Operand resolved to a prior selection (e.g. ARRAY_AGG) which produces a native PG array. return String.format( "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", resolvedSelection.get(), ARRAY_DIMENSION); } - // Raw jsonb field access — use jsonb_array_length which operates on jsonb arrays directly. - String parsedExpression = operand.accept(selectTypeExpressionVisitor); + PostgresFieldIdentifierExpressionVisitor jsonbVisitor = + new PostgresFieldIdentifierExpressionVisitor(getPostgresQueryParser()); + String parsedExpression = operand.accept(jsonbVisitor); return String.format("COALESCE( jsonb_array_length( %s ), 0 )", parsedExpression); } From b7b74390ef39946f388b39056208a3c32aa15ef2 Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Fri, 19 Jun 2026 23:20:11 +0530 Subject: [PATCH 06/12] fix: use ARRAY_LENGTH for flat collections in LENGTH operator For flat collections where fields are native PG arrays (e.g. TEXT[]), use ARRAY_LENGTH. For nested collections with jsonb document fields, use jsonb_array_length. Both wrapped with COALESCE to return 0 for NULL/missing. Co-Authored-By: Claude Opus 4.6 --- .../v1/vistors/PostgresFunctionExpressionVisitor.java | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java index 9cb85e29c..4008028e1 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java @@ -6,6 +6,7 @@ import java.util.stream.Collectors; import lombok.NoArgsConstructor; import org.apache.commons.lang3.StringUtils; +import org.hypertrace.core.documentstore.DocumentType; import org.hypertrace.core.documentstore.expression.impl.FunctionExpression; import org.hypertrace.core.documentstore.expression.operators.FunctionOperator; import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression; @@ -92,9 +93,13 @@ private String buildLengthExpression(final SelectTypeExpression operand) { return String.format( "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", resolvedSelection.get(), ARRAY_DIMENSION); } - PostgresFieldIdentifierExpressionVisitor jsonbVisitor = + PostgresFieldIdentifierExpressionVisitor fieldVisitor = new PostgresFieldIdentifierExpressionVisitor(getPostgresQueryParser()); - String parsedExpression = operand.accept(jsonbVisitor); + String parsedExpression = operand.accept(fieldVisitor); + if (getPostgresQueryParser().getPgColTransformer().getDocumentType() == DocumentType.FLAT) { + return String.format( + "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", parsedExpression, ARRAY_DIMENSION); + } return String.format("COALESCE( jsonb_array_length( %s ), 0 )", parsedExpression); } From be9514adf43bdfd7d4f8ba3125b06d6ea7496663 Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Fri, 19 Jun 2026 23:26:55 +0530 Subject: [PATCH 07/12] refactor: wrap test body in try-finally for deterministic cleanup Ensures deleteCollection runs even if assertions throw. Co-Authored-By: Claude Opus 4.6 --- .../documentstore/DocStoreQueryV1Test.java | 73 ++++++++++--------- 1 file changed, 38 insertions(+), 35 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java index a1346bc08..4bc3397be 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java @@ -928,44 +928,47 @@ public void testSortByListSizeWithMissingField(String dataStoreName) throws IOEx datastore.createCollection(collectionName, null); Collection collection = datastore.getCollection(collectionName); - collection.upsert( - new SingleValueKey(TENANT_ID, "three"), - new JSONDocument("{\"item\":\"three\",\"tags\":[\"a\",\"b\",\"c\"]}")); - collection.upsert( - new SingleValueKey(TENANT_ID, "one"), - new JSONDocument("{\"item\":\"one\",\"tags\":[\"x\"]}")); - // Document intentionally missing the "tags" field; LENGTH must resolve to 0 instead of failing - collection.upsert( - new SingleValueKey(TENANT_ID, "none"), new JSONDocument("{\"item\":\"none\"}")); + try { + collection.upsert( + new SingleValueKey(TENANT_ID, "three"), + new JSONDocument("{\"item\":\"three\",\"tags\":[\"a\",\"b\",\"c\"]}")); + collection.upsert( + new SingleValueKey(TENANT_ID, "one"), + new JSONDocument("{\"item\":\"one\",\"tags\":[\"x\"]}")); + // Document intentionally missing the "tags" field; LENGTH must resolve to 0 instead of + // failing + collection.upsert( + new SingleValueKey(TENANT_ID, "none"), new JSONDocument("{\"item\":\"none\"}")); - Query query = - Query.builder() - .addSelection(IdentifierExpression.of("item")) - .addSelection( - FunctionExpression.builder() - .operator(LENGTH) - .operand(IdentifierExpression.of("tags")) - .build(), - "tag_count") - .addSort(IdentifierExpression.of("tag_count"), ASC) - .build(); + Query query = + Query.builder() + .addSelection(IdentifierExpression.of("item")) + .addSelection( + FunctionExpression.builder() + .operator(LENGTH) + .operand(IdentifierExpression.of("tags")) + .build(), + "tag_count") + .addSort(IdentifierExpression.of("tag_count"), ASC) + .build(); + + Iterator resultDocs = collection.aggregate(query); + List> results = new ArrayList<>(); + while (resultDocs.hasNext()) { + results.add(Utils.convertDocumentToMap(resultDocs.next())); + } - Iterator resultDocs = collection.aggregate(query); - List> results = new ArrayList<>(); - while (resultDocs.hasNext()) { - results.add(Utils.convertDocumentToMap(resultDocs.next())); + assertEquals(3, results.size()); + // Document without the "tags" field counts as 0 and sorts first in ascending order + assertEquals("none", results.get(0).get("item")); + assertEquals(0, ((Number) results.get(0).get("tag_count")).intValue()); + assertEquals("one", results.get(1).get("item")); + assertEquals(1, ((Number) results.get(1).get("tag_count")).intValue()); + assertEquals("three", results.get(2).get("item")); + assertEquals(3, ((Number) results.get(2).get("tag_count")).intValue()); + } finally { + datastore.deleteCollection(collectionName); } - - assertEquals(3, results.size()); - // Document without the "tags" field counts as 0 and sorts first in ascending order - assertEquals("none", results.get(0).get("item")); - assertEquals(0, ((Number) results.get(0).get("tag_count")).intValue()); - assertEquals("one", results.get(1).get("item")); - assertEquals(1, ((Number) results.get(1).get("tag_count")).intValue()); - assertEquals("three", results.get(2).get("item")); - assertEquals(3, ((Number) results.get(2).get("tag_count")).intValue()); - - datastore.deleteCollection(collectionName); } @ParameterizedTest From bc24f4c784b74f83c0d0af48747433de11de4a8a Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Fri, 19 Jun 2026 23:34:16 +0530 Subject: [PATCH 08/12] fix: handle JSON null in jsonb LENGTH using jsonb_typeof guard COALESCE only handles SQL NULL (missing key). When a field has an explicit JSON null value, jsonb_array_length errors. Use the same CASE WHEN jsonb_typeof = 'array' pattern used elsewhere in the codebase to safely return 0 for missing, null, and non-array fields. Co-Authored-By: Claude Opus 4.6 --- .../query/v1/vistors/PostgresFunctionExpressionVisitor.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java index 4008028e1..f59ee9d01 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java @@ -100,7 +100,10 @@ private String buildLengthExpression(final SelectTypeExpression operand) { return String.format( "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", parsedExpression, ARRAY_DIMENSION); } - return String.format("COALESCE( jsonb_array_length( %s ), 0 )", parsedExpression); + return String.format( + "jsonb_array_length( CASE WHEN jsonb_typeof( %s ) = 'array' THEN %s" + + " ELSE '[]'::jsonb END )", + parsedExpression, parsedExpression); } private String getParsedExpression(final SelectTypeExpression expression) { From df178fa3b9c37577865d1e482d8ff6a15f74cf87 Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Wed, 24 Jun 2026 10:43:20 +0530 Subject: [PATCH 09/12] test: add integration test for LENGTH on jsonb field in flat collection Tests LENGTH on JsonIdentifierExpression.of("props", "colors") in a flat Postgres collection where props is a JSONB column. This should fail with current impl (ARRAY_LENGTH on jsonb) and pass after fix. Co-Authored-By: Claude Opus 4.6 --- .../documentstore/DocStoreQueryV1Test.java | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java index 4bc3397be..302c5e269 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreQueryV1Test.java @@ -971,6 +971,46 @@ public void testSortByListSizeWithMissingField(String dataStoreName) throws IOEx } } + @ParameterizedTest + @ArgumentsSource(PostgresProvider.class) + public void testSortByJsonbArrayLengthInFlatCollection(String dataStoreName) throws IOException { + Collection flatCollection = getFlatCollection(dataStoreName); + + // props.colors is a jsonb array inside a JSONB column in the flat collection. + // Test data: id=1 ["Blue","Green"], id=3 ["Black"], id=5 ["Orange","Blue"], + // id=7 [] (empty), id=2,4,6,8,9,10 props=NULL. + Query query = + Query.builder() + .addSelection(IdentifierExpression.of("id")) + .addSelection( + FunctionExpression.builder() + .operator(LENGTH) + .operand(JsonIdentifierExpression.of("props", "colors")) + .build(), + "color_count") + .addSort(IdentifierExpression.of("color_count"), DESC) + .addSort(IdentifierExpression.of("id"), ASC) + .build(); + + Iterator resultDocs = flatCollection.aggregate(query); + List> results = new ArrayList<>(); + while (resultDocs.hasNext()) { + results.add(Utils.convertDocumentToMap(resultDocs.next())); + } + + assertEquals(10, results.size()); + // id=1 and id=5 have 2 colors + assertEquals("1", results.get(0).get("id")); + assertEquals(2, ((Number) results.get(0).get("color_count")).intValue()); + assertEquals("5", results.get(1).get("id")); + assertEquals(2, ((Number) results.get(1).get("color_count")).intValue()); + // id=3 has 1 color + assertEquals("3", results.get(2).get("id")); + assertEquals(1, ((Number) results.get(2).get("color_count")).intValue()); + // Remaining have 0 (empty array or NULL props) + assertEquals(0, ((Number) results.get(3).get("color_count")).intValue()); + } + @ParameterizedTest @ArgumentsSource(AllProvider.class) public void testAggregateWithFunctionalLeftHandSideFilter(final String dataStoreName) From e75d0455ee131111ec05ba1a3ef07ab074ce92d6 Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Wed, 24 Jun 2026 10:51:13 +0530 Subject: [PATCH 10/12] fix: use FieldToPgColumn to determine array type instead of DocumentType Check whether the field resolves to a direct column (native PG array) or a jsonb path (jsonb array) by inspecting FieldToPgColumn.transformedField. This correctly handles flat collections with JSONB columns where fields inside the JSONB column need jsonb_array_length, not ARRAY_LENGTH. Co-Authored-By: Claude Opus 4.6 --- .../PostgresFunctionExpressionVisitor.java | 23 +++++++++++++++---- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java index f59ee9d01..fc2a9a28b 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java @@ -6,11 +6,12 @@ import java.util.stream.Collectors; import lombok.NoArgsConstructor; import org.apache.commons.lang3.StringUtils; -import org.hypertrace.core.documentstore.DocumentType; import org.hypertrace.core.documentstore.expression.impl.FunctionExpression; +import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression; import org.hypertrace.core.documentstore.expression.operators.FunctionOperator; import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; +import org.hypertrace.core.documentstore.postgres.query.v1.transformer.FieldToPgColumn; @NoArgsConstructor public class PostgresFunctionExpressionVisitor extends PostgresSelectTypeExpressionVisitor { @@ -93,19 +94,31 @@ private String buildLengthExpression(final SelectTypeExpression operand) { return String.format( "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", resolvedSelection.get(), ARRAY_DIMENSION); } - PostgresFieldIdentifierExpressionVisitor fieldVisitor = - new PostgresFieldIdentifierExpressionVisitor(getPostgresQueryParser()); - String parsedExpression = operand.accept(fieldVisitor); - if (getPostgresQueryParser().getPgColTransformer().getDocumentType() == DocumentType.FLAT) { + if (isDirectColumn(operand)) { + PostgresFieldIdentifierExpressionVisitor fieldVisitor = + new PostgresFieldIdentifierExpressionVisitor(getPostgresQueryParser()); + String parsedExpression = operand.accept(fieldVisitor); return String.format( "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", parsedExpression, ARRAY_DIMENSION); } + PostgresFieldIdentifierExpressionVisitor fieldVisitor = + new PostgresFieldIdentifierExpressionVisitor(getPostgresQueryParser()); + String parsedExpression = operand.accept(fieldVisitor); return String.format( "jsonb_array_length( CASE WHEN jsonb_typeof( %s ) = 'array' THEN %s" + " ELSE '[]'::jsonb END )", parsedExpression, parsedExpression); } + private boolean isDirectColumn(final SelectTypeExpression operand) { + if (operand instanceof IdentifierExpression) { + FieldToPgColumn fieldToPgColumn = + getPostgresQueryParser().transformField((IdentifierExpression) operand); + return fieldToPgColumn.getTransformedField() == null; + } + return false; + } + private String getParsedExpression(final SelectTypeExpression expression) { Optional identifier = Optional.ofNullable(expression.accept(identifierExpressionVisitor)); From 5fb58fad1e9b213c39822ac249a7d60355dead33 Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Thu, 25 Jun 2026 13:24:40 +0530 Subject: [PATCH 11/12] refactor: move array length logic into PostgresColTransformer Add buildArrayLengthExpression to PostgresColTransformer so storage-layout decisions (native PG array vs jsonb) stay in the transformer layer. The function visitor now delegates to the transformer instead of inspecting FieldToPgColumn internals directly. Co-Authored-By: Claude Opus 4.8 --- .../FlatPostgresFieldTransformer.java | 12 +++++++++++ .../NestedPostgresColTransformer.java | 9 ++++++++ .../transformer/PostgresColTransformer.java | 10 +++++++++ .../PostgresFunctionExpressionVisitor.java | 21 +++++-------------- 4 files changed, 36 insertions(+), 16 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FlatPostgresFieldTransformer.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FlatPostgresFieldTransformer.java index 9c6a4dfe2..a8ea2b3f1 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FlatPostgresFieldTransformer.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FlatPostgresFieldTransformer.java @@ -103,6 +103,18 @@ public String buildFieldAccessorWithoutCast(FieldToPgColumn fieldToPgColumn) { .toString(); } + @Override + public String buildArrayLengthExpression(FieldToPgColumn fieldToPgColumn) { + String fieldAccessor = buildFieldAccessorWithoutCast(fieldToPgColumn); + if (fieldToPgColumn.getTransformedField() == null) { + return String.format("COALESCE( ARRAY_LENGTH( %s, 1 ), 0 )", fieldAccessor); + } + return String.format( + "jsonb_array_length( CASE WHEN jsonb_typeof( %s ) = 'array' THEN %s" + + " ELSE '[]'::jsonb END )", + fieldAccessor, fieldAccessor); + } + @Override public DocumentType getDocumentType() { return DocumentType.FLAT; diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/NestedPostgresColTransformer.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/NestedPostgresColTransformer.java index 3eee307c5..6d425dc63 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/NestedPostgresColTransformer.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/NestedPostgresColTransformer.java @@ -103,6 +103,15 @@ public String buildFieldAccessorWithoutCast(FieldToPgColumn fieldToPgColumn) { .toString(); } + @Override + public String buildArrayLengthExpression(FieldToPgColumn fieldToPgColumn) { + String fieldAccessor = buildFieldAccessorWithoutCast(fieldToPgColumn); + return String.format( + "jsonb_array_length( CASE WHEN jsonb_typeof( %s ) = 'array' THEN %s" + + " ELSE '[]'::jsonb END )", + fieldAccessor, fieldAccessor); + } + @Override public DocumentType getDocumentType() { return DocumentType.NESTED; diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresColTransformer.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresColTransformer.java index 4d842ff35..5fa76192d 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresColTransformer.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/PostgresColTransformer.java @@ -38,6 +38,16 @@ public interface PostgresColTransformer { */ String buildFieldAccessorWithoutCast(FieldToPgColumn fieldToPgColumn); + /** + * Builds a SQL expression that computes the length of an array field, returning 0 for + * NULL/missing/empty arrays. Implementations choose between ARRAY_LENGTH (native PG arrays) and + * jsonb_array_length (JSONB arrays) based on the field's storage type. + * + * @param fieldToPgColumn The result of field transformation + * @return SQL expression computing the array length with NULL-safe handling + */ + String buildArrayLengthExpression(FieldToPgColumn fieldToPgColumn); + /** * Returns the kind of document this transformer is handling - Flat vs nested * diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java index fc2a9a28b..5c38fe8b7 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java @@ -11,7 +11,6 @@ import org.hypertrace.core.documentstore.expression.operators.FunctionOperator; import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; -import org.hypertrace.core.documentstore.postgres.query.v1.transformer.FieldToPgColumn; @NoArgsConstructor public class PostgresFunctionExpressionVisitor extends PostgresSelectTypeExpressionVisitor { @@ -94,12 +93,11 @@ private String buildLengthExpression(final SelectTypeExpression operand) { return String.format( "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", resolvedSelection.get(), ARRAY_DIMENSION); } - if (isDirectColumn(operand)) { - PostgresFieldIdentifierExpressionVisitor fieldVisitor = - new PostgresFieldIdentifierExpressionVisitor(getPostgresQueryParser()); - String parsedExpression = operand.accept(fieldVisitor); - return String.format( - "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", parsedExpression, ARRAY_DIMENSION); + if (operand instanceof IdentifierExpression) { + PostgresQueryParser parser = getPostgresQueryParser(); + return parser + .getPgColTransformer() + .buildArrayLengthExpression(parser.transformField((IdentifierExpression) operand)); } PostgresFieldIdentifierExpressionVisitor fieldVisitor = new PostgresFieldIdentifierExpressionVisitor(getPostgresQueryParser()); @@ -110,15 +108,6 @@ private String buildLengthExpression(final SelectTypeExpression operand) { parsedExpression, parsedExpression); } - private boolean isDirectColumn(final SelectTypeExpression operand) { - if (operand instanceof IdentifierExpression) { - FieldToPgColumn fieldToPgColumn = - getPostgresQueryParser().transformField((IdentifierExpression) operand); - return fieldToPgColumn.getTransformedField() == null; - } - return false; - } - private String getParsedExpression(final SelectTypeExpression expression) { Optional identifier = Optional.ofNullable(expression.accept(identifierExpressionVisitor)); From 08650aba909fcbd6b7e3512a0282b3d03d9c64b5 Mon Sep 17 00:00:00 2001 From: sarthak77 Date: Fri, 26 Jun 2026 11:47:12 +0530 Subject: [PATCH 12/12] refactor: address review feedback on array length handling - Rename visitor's buildLengthExpression -> buildArrayLengthExpression for consistency with the transformer interface method. - Extract the COALESCE(ARRAY_LENGTH(...)) and jsonb_array_length(CASE WHEN jsonb_typeof...) templates into PostgresUtils.prepareArrayLength / prepareJsonbArrayLength to remove duplication. - Drop the dead non-identifier fallback branch; LENGTH operands are always an alias or a field identifier. - Replace Optional plumbing with a plain null check in the LENGTH path. Co-Authored-By: Claude Opus 4.8 --- .../FlatPostgresFieldTransformer.java | 12 +++--- .../NestedPostgresColTransformer.java | 7 +--- .../PostgresFunctionExpressionVisitor.java | 39 ++++++++----------- .../postgres/utils/PostgresUtils.java | 19 +++++++++ 4 files changed, 43 insertions(+), 34 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FlatPostgresFieldTransformer.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FlatPostgresFieldTransformer.java index a8ea2b3f1..0cf619e5a 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FlatPostgresFieldTransformer.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/FlatPostgresFieldTransformer.java @@ -106,13 +106,11 @@ public String buildFieldAccessorWithoutCast(FieldToPgColumn fieldToPgColumn) { @Override public String buildArrayLengthExpression(FieldToPgColumn fieldToPgColumn) { String fieldAccessor = buildFieldAccessorWithoutCast(fieldToPgColumn); - if (fieldToPgColumn.getTransformedField() == null) { - return String.format("COALESCE( ARRAY_LENGTH( %s, 1 ), 0 )", fieldAccessor); - } - return String.format( - "jsonb_array_length( CASE WHEN jsonb_typeof( %s ) = 'array' THEN %s" - + " ELSE '[]'::jsonb END )", - fieldAccessor, fieldAccessor); + // A direct column (no JSON path) is a native PG array; otherwise it is a JSONB array nested in + // a JSONB column. + return fieldToPgColumn.getTransformedField() == null + ? PostgresUtils.prepareArrayLength(fieldAccessor) + : PostgresUtils.prepareJsonbArrayLength(fieldAccessor); } @Override diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/NestedPostgresColTransformer.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/NestedPostgresColTransformer.java index 6d425dc63..f37d53475 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/NestedPostgresColTransformer.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/transformer/NestedPostgresColTransformer.java @@ -105,11 +105,8 @@ public String buildFieldAccessorWithoutCast(FieldToPgColumn fieldToPgColumn) { @Override public String buildArrayLengthExpression(FieldToPgColumn fieldToPgColumn) { - String fieldAccessor = buildFieldAccessorWithoutCast(fieldToPgColumn); - return String.format( - "jsonb_array_length( CASE WHEN jsonb_typeof( %s ) = 'array' THEN %s" - + " ELSE '[]'::jsonb END )", - fieldAccessor, fieldAccessor); + // Nested collections store every field inside the JSONB document column. + return PostgresUtils.prepareJsonbArrayLength(buildFieldAccessorWithoutCast(fieldToPgColumn)); } @Override diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java index 5c38fe8b7..b20821c70 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFunctionExpressionVisitor.java @@ -11,11 +11,11 @@ import org.hypertrace.core.documentstore.expression.operators.FunctionOperator; import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; +import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils; @NoArgsConstructor public class PostgresFunctionExpressionVisitor extends PostgresSelectTypeExpressionVisitor { - private static final int ARRAY_DIMENSION = 1; private PostgresIdentifierExpressionVisitor identifierExpressionVisitor; private PostgresSelectTypeExpressionVisitor selectTypeExpressionVisitor; @@ -51,7 +51,7 @@ public String visit(final FunctionExpression expression) { if (numArgs == 1) { if (expression.getOperator().equals(FunctionOperator.LENGTH)) { - return buildLengthExpression(expression.getOperands().get(0)); + return buildArrayLengthExpression(expression.getOperands().get(0)); } String parsedExpression = getParsedExpression(expression.getOperands().get(0)); return String.format("%s( %s )", expression.getOperator(), parsedExpression); @@ -85,27 +85,22 @@ private Collector getCollectorForFunctionOperator(FunctionOperator operator) { String.format("Query operation:%s not supported", operator)); } - private String buildLengthExpression(final SelectTypeExpression operand) { - Optional identifier = Optional.ofNullable(operand.accept(identifierExpressionVisitor)); - Optional resolvedSelection = - identifier.map(v -> getPostgresQueryParser().getPgSelections().get(v)); - if (resolvedSelection.isPresent()) { - return String.format( - "COALESCE( ARRAY_LENGTH( %s, %s ), 0 )", resolvedSelection.get(), ARRAY_DIMENSION); - } - if (operand instanceof IdentifierExpression) { - PostgresQueryParser parser = getPostgresQueryParser(); - return parser - .getPgColTransformer() - .buildArrayLengthExpression(parser.transformField((IdentifierExpression) operand)); + private String buildArrayLengthExpression(final SelectTypeExpression operand) { + PostgresQueryParser parser = getPostgresQueryParser(); + + // The operand is either a user-defined alias from a prior selection (e.g. ARRAY_AGG, which + // produces a native PG array) or a field identifier. Aliases are resolved against pgSelections. + String identifier = operand.accept(identifierExpressionVisitor); + String resolvedSelection = identifier != null ? parser.getPgSelections().get(identifier) : null; + if (resolvedSelection != null) { + return PostgresUtils.prepareArrayLength(resolvedSelection); } - PostgresFieldIdentifierExpressionVisitor fieldVisitor = - new PostgresFieldIdentifierExpressionVisitor(getPostgresQueryParser()); - String parsedExpression = operand.accept(fieldVisitor); - return String.format( - "jsonb_array_length( CASE WHEN jsonb_typeof( %s ) = 'array' THEN %s" - + " ELSE '[]'::jsonb END )", - parsedExpression, parsedExpression); + + // A plain field identifier — let the transformer decide between ARRAY_LENGTH (native array) and + // jsonb_array_length (JSONB array) based on the field's storage layout. + return parser + .getPgColTransformer() + .buildArrayLengthExpression(parser.transformField((IdentifierExpression) operand)); } private String getParsedExpression(final SelectTypeExpression expression) { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/utils/PostgresUtils.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/utils/PostgresUtils.java index 5e8e64cf7..2d6860d03 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/utils/PostgresUtils.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/utils/PostgresUtils.java @@ -131,6 +131,25 @@ public static String prepareCast(String field, Type type) { } } + /** + * Builds a NULL-safe length expression for a native PG array (e.g. {@code TEXT[]}). Returns 0 for + * a NULL or empty array. + */ + public static String prepareArrayLength(final String arrayAccessor) { + return String.format("COALESCE( ARRAY_LENGTH( %s, 1 ), 0 )", arrayAccessor); + } + + /** + * Builds a length expression for a JSONB value. The {@code jsonb_typeof} guard makes absent, JSON + * {@code null}, and non-array values resolve to 0 instead of erroring inside {@code + * jsonb_array_length}. + */ + public static String prepareJsonbArrayLength(final String jsonbAccessor) { + return String.format( + "jsonb_array_length( CASE WHEN jsonb_typeof( %s ) = 'array' THEN %s ELSE '[]'::jsonb END )", + jsonbAccessor, jsonbAccessor); + } + public static String prepareCastForFieldAccessor(String field, Object value) { String fmt = "CAST (%s AS %s)"; Type type = getType(value);