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..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 @@ -919,6 +919,58 @@ public void testAggregateWithMultipleGroupingLevels(String dataStoreName) throws testCountApi(dataStoreName, query, "query/multi_level_grouping_response.json"); } + @ParameterizedTest + @ArgumentsSource(AllProvider.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); + + 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(); + + 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); + } + } + @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..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 @@ -69,6 +69,12 @@ 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..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 @@ -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; @@ -49,10 +50,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) - ? String.format("ARRAY_LENGTH( %s, %s )", parsedExpression, ARRAY_DIMENSION) - : String.format("%s( %s )", expression.getOperator(), parsedExpression); + return String.format("%s( %s )", expression.getOperator(), parsedExpression); } Collector collector = @@ -83,6 +85,27 @@ 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); + } + PostgresFieldIdentifierExpressionVisitor fieldVisitor = + new PostgresFieldIdentifierExpressionVisitor(getPostgresQueryParser()); + 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( + "jsonb_array_length( CASE WHEN jsonb_typeof( %s ) = 'array' THEN %s" + + " ELSE '[]'::jsonb END )", + parsedExpression, parsedExpression); + } + private String getParsedExpression(final SelectTypeExpression expression) { Optional identifier = Optional.ofNullable(expression.accept(identifierExpressionVisitor)); 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", + [] + ] + } } } }