diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/ArrayFiltersQueryIntegrationTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/ArrayFiltersQueryIntegrationTest.java index 318bafa67..e6f066a51 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/ArrayFiltersQueryIntegrationTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/ArrayFiltersQueryIntegrationTest.java @@ -2,7 +2,9 @@ import static org.hypertrace.core.documentstore.expression.impl.LogicalExpression.and; import static org.hypertrace.core.documentstore.expression.impl.LogicalExpression.not; +import static org.hypertrace.core.documentstore.expression.operators.ArrayOperator.ALL; import static org.hypertrace.core.documentstore.expression.operators.ArrayOperator.ANY; +import static org.hypertrace.core.documentstore.expression.operators.ArrayOperator.ONE; import static org.hypertrace.core.documentstore.model.config.DatabaseType.MONGO; import static org.hypertrace.core.documentstore.model.config.DatabaseType.POSTGRES; import static org.hypertrace.core.documentstore.utils.Utils.MONGO_STORE; @@ -11,6 +13,7 @@ import com.google.common.io.Resources; import java.io.IOException; import java.util.Iterator; +import java.util.List; import java.util.Map; import java.util.Spliterator; import java.util.Spliterators; @@ -344,6 +347,84 @@ void getDocumentsWithEnvironmentIdsSubsetOfGivenList(final String dataStoreName) JSONAssert.assertEquals(expected, actual, JSONCompareMode.LENIENT); } + /** + * Tests MATCH_ALL semantics: documents whose array attribute contains every value specified in + * the filter. + */ + @ParameterizedTest + @ArgumentsSource(AllProvider.class) + void getDocumentsContainingAllGivenValues(final String dataStoreName) + throws JSONException, IOException { + final String testCollectionName = "array_match_test"; + final Datastore datastore = datastoreMap.get(dataStoreName); + final Map testDocuments = + Utils.buildDocumentsFromResource("query/array_operators/array_match_test.json"); + datastore.deleteCollection(testCollectionName); + datastore.createCollection(testCollectionName, null); + final Collection collection = datastore.getCollection(testCollectionName); + collection.bulkUpsert(testDocuments); + + final Query query = + Query.builder() + .setFilter( + ArrayRelationalFilterExpression.builder() + .operator(ALL) + .filter( + RelationalExpression.of( + IdentifierExpression.of("tags"), + RelationalOperator.IN, + ConstantExpression.ofStrings(List.of("red", "blue")))) + .build()) + .build(); + + final Iterator documents = collection.aggregate(query); + final String expected = readResource("array_match_all_result.json"); + final String actual = iteratorToJson(documents); + + datastore.deleteCollection(testCollectionName); + + JSONAssert.assertEquals(expected, actual, JSONCompareMode.LENIENT); + } + + /** + * Tests MATCH_ONE semantics: documents whose array attribute has exactly one element, and that + * element is one of the values specified in the filter. + */ + @ParameterizedTest + @ArgumentsSource(AllProvider.class) + void getDocumentsWithExactlyOneElementMatchingGivenValues(final String dataStoreName) + throws JSONException, IOException { + final String testCollectionName = "array_match_test"; + final Datastore datastore = datastoreMap.get(dataStoreName); + final Map testDocuments = + Utils.buildDocumentsFromResource("query/array_operators/array_match_test.json"); + datastore.deleteCollection(testCollectionName); + datastore.createCollection(testCollectionName, null); + final Collection collection = datastore.getCollection(testCollectionName); + collection.bulkUpsert(testDocuments); + + final Query query = + Query.builder() + .setFilter( + ArrayRelationalFilterExpression.builder() + .operator(ONE) + .filter( + RelationalExpression.of( + IdentifierExpression.of("tags"), + RelationalOperator.IN, + ConstantExpression.ofStrings(List.of("red", "blue")))) + .build()) + .build(); + + final Iterator documents = collection.aggregate(query); + final String expected = readResource("array_match_one_result.json"); + final String actual = iteratorToJson(documents); + + datastore.deleteCollection(testCollectionName); + + JSONAssert.assertEquals(expected, actual, JSONCompareMode.LENIENT); + } + private String readResource(final String fileName) { try { return new String( diff --git a/document-store/src/integrationTest/resources/query/array_operators/array_match_all_result.json b/document-store/src/integrationTest/resources/query/array_operators/array_match_all_result.json new file mode 100644 index 000000000..d9fd9221a --- /dev/null +++ b/document-store/src/integrationTest/resources/query/array_operators/array_match_all_result.json @@ -0,0 +1,10 @@ +[ + { + "name": "Document A", + "tags": ["red", "blue"] + }, + { + "name": "Document B", + "tags": ["red", "blue", "green"] + } +] diff --git a/document-store/src/integrationTest/resources/query/array_operators/array_match_one_result.json b/document-store/src/integrationTest/resources/query/array_operators/array_match_one_result.json new file mode 100644 index 000000000..2f5de65f9 --- /dev/null +++ b/document-store/src/integrationTest/resources/query/array_operators/array_match_one_result.json @@ -0,0 +1,6 @@ +[ + { + "name": "Document C", + "tags": ["red"] + } +] diff --git a/document-store/src/integrationTest/resources/query/array_operators/array_match_test.json b/document-store/src/integrationTest/resources/query/array_operators/array_match_test.json new file mode 100644 index 000000000..29b1a1a07 --- /dev/null +++ b/document-store/src/integrationTest/resources/query/array_operators/array_match_test.json @@ -0,0 +1,31 @@ +[ + { + "_id": 1, + "name": "Document A", + "tags": ["red", "blue"] + }, + { + "_id": 2, + "name": "Document B", + "tags": ["red", "blue", "green"] + }, + { + "_id": 3, + "name": "Document C", + "tags": ["red"] + }, + { + "_id": 4, + "name": "Document D", + "tags": ["yellow"] + }, + { + "_id": 5, + "name": "Document E", + "tags": [] + }, + { + "_id": 6, + "name": "Document F" + } +] diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/expression/operators/ArrayOperator.java b/document-store/src/main/java/org/hypertrace/core/documentstore/expression/operators/ArrayOperator.java index b018032e1..319741858 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/expression/operators/ArrayOperator.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/expression/operators/ArrayOperator.java @@ -2,5 +2,9 @@ public enum ArrayOperator { ANY, - // Can support ALL and NONE later + // Array attribute must contain every value specified in the filter + ALL, + // Array attribute must contain exactly one element, and that element must be one of the values + // specified in the filter + ONE, } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoArrayFilterParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoArrayFilterParser.java index cab7312c6..5269ac4e7 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoArrayFilterParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoArrayFilterParser.java @@ -6,10 +6,14 @@ import static org.hypertrace.core.documentstore.mongo.query.parser.filter.MongoStandardExprRelationalFilterParser.EXPR; import com.google.common.collect.Maps; +import java.util.List; import java.util.Map; import java.util.Optional; import org.hypertrace.core.documentstore.expression.impl.ArrayFilterExpression; +import org.hypertrace.core.documentstore.expression.impl.ConstantExpression; +import org.hypertrace.core.documentstore.expression.impl.RelationalExpression; import org.hypertrace.core.documentstore.expression.operators.ArrayOperator; +import org.hypertrace.core.documentstore.expression.type.FilterTypeExpression; import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression; import org.hypertrace.core.documentstore.mongo.MongoUtils; import org.hypertrace.core.documentstore.mongo.query.parser.filter.MongoRelationalFilterParserFactory.MongoRelationalFilterContext; @@ -21,6 +25,12 @@ class MongoArrayFilterParser { private static final String IF_NULL = "$ifNull"; private static final String AS = "as"; private static final String IN = "in"; + private static final String SET_IS_SUBSET = "$setIsSubset"; + private static final String SIZE = "$size"; + private static final String EQ = "$eq"; + private static final String IN_OPERATOR = "$in"; + private static final String ARRAY_ELEM_AT = "$arrayElemAt"; + private static final String AND = "$and"; private static final Map OPERATOR_MAP = Maps.immutableEnumMap(Map.ofEntries(entry(ANY, ANY_ELEMENT_TRUE))); @@ -39,6 +49,17 @@ class MongoArrayFilterParser { } Map parse(final ArrayFilterExpression arrayFilterExpression) { + switch (arrayFilterExpression.getOperator()) { + case ALL: + return parseAllOperator(arrayFilterExpression); + case ONE: + return parseOneOperator(arrayFilterExpression); + default: + return parseAnyOperator(arrayFilterExpression); + } + } + + private Map parseAnyOperator(final ArrayFilterExpression arrayFilterExpression) { final String operator = Optional.ofNullable(OPERATOR_MAP.get(arrayFilterExpression.getOperator())) .orElseThrow( @@ -103,9 +124,84 @@ Map parse(final ArrayFilterExpression arrayFilterExpression) { entry(INPUT, Map.of(IF_NULL, new Object[] {mapInput, new Object[0]})), entry(AS, alias), entry(IN, filter)))); + return wrapInExprIfNeeded(arrayFilter); + } + + /* + { + "$expr": { + "$setIsSubset": [ + ["Blue", "Green"], + { "$ifNull": ["$colors", []] } + ] + } + } + */ + private Map parseAllOperator(final ArrayFilterExpression arrayFilterExpression) { + final Object mapInput = getDollarPrefixedArraySource(arrayFilterExpression); + final List values = getFilterValues(arrayFilterExpression); + + final Map setIsSubset = + Map.of( + SET_IS_SUBSET, + List.of(values, Map.of(IF_NULL, new Object[] {mapInput, new Object[0]}))); + return wrapInExprIfNeeded(setIsSubset); + } + + /* + { + "$expr": { + "$and": [ + { "$eq": [{ "$size": { "$ifNull": ["$colors", []] } }, 1] }, + { "$in": [{ "$arrayElemAt": [{ "$ifNull": ["$colors", []] }, 0] }, ["Blue", "Green"]] } + ] + } + } + */ + private Map parseOneOperator(final ArrayFilterExpression arrayFilterExpression) { + final Object mapInput = getDollarPrefixedArraySource(arrayFilterExpression); + final List values = getFilterValues(arrayFilterExpression); + final Map arrayWithDefault = + Map.of(IF_NULL, new Object[] {mapInput, new Object[0]}); + + final Map sizeIsOne = Map.of(EQ, List.of(Map.of(SIZE, arrayWithDefault), 1)); + final Map firstElementMatches = + Map.of(IN_OPERATOR, List.of(Map.of(ARRAY_ELEM_AT, List.of(arrayWithDefault, 0)), values)); + + return wrapInExprIfNeeded(Map.of(AND, List.of(sizeIsOne, firstElementMatches))); + } + + private String getDollarPrefixedArraySource(final ArrayFilterExpression arrayFilterExpression) { + final MongoSelectTypeExpressionParser wrappingParser = + new MongoDollarPrefixingIdempotentParser(relationalFilterContext.lhsParser()); + return arrayFilterExpression.getArraySource().accept(wrappingParser); + } + + private List getFilterValues(final ArrayFilterExpression arrayFilterExpression) { + final FilterTypeExpression filter = arrayFilterExpression.getFilter(); + if (!(filter instanceof RelationalExpression)) { + throw new UnsupportedOperationException( + "Array operator " + + arrayFilterExpression.getOperator() + + " only supports a relational filter with a constant list of values, got: " + + filter); + } + + final SelectTypeExpression rhs = ((RelationalExpression) filter).getRhs(); + if (!(rhs instanceof ConstantExpression)) { + throw new UnsupportedOperationException( + "Array operator " + + arrayFilterExpression.getOperator() + + " requires a constant list of values, got: " + + rhs); + } + + final Object value = ((ConstantExpression) rhs).getValue(); + return value instanceof List ? (List) value : List.of(value); + } + + private Map wrapInExprIfNeeded(final Map filter) { // If already wrapped inside `$expr` avoid wrapping again - return INSIDE_EXPR.equals(relationalFilterContext.location()) - ? arrayFilter - : Map.of(EXPR, arrayFilter); + return INSIDE_EXPR.equals(relationalFilterContext.location()) ? filter : Map.of(EXPR, filter); } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java index cab7713b7..8de9c1fe8 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java @@ -8,6 +8,9 @@ import static org.hypertrace.core.documentstore.postgres.utils.PostgresUtils.encodeAliasForNestedField; import static org.hypertrace.core.documentstore.postgres.utils.PostgresUtils.prepareParsedNonCompositeFilter; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.List; import java.util.Optional; import java.util.stream.Collector; import java.util.stream.Collectors; @@ -15,6 +18,7 @@ import org.apache.commons.lang3.StringUtils; import org.hypertrace.core.documentstore.DocumentType; import org.hypertrace.core.documentstore.Key; +import org.hypertrace.core.documentstore.expression.impl.ArrayFilterExpression; import org.hypertrace.core.documentstore.expression.impl.ArrayRelationalFilterExpression; import org.hypertrace.core.documentstore.expression.impl.ConstantExpression; import org.hypertrace.core.documentstore.expression.impl.DocumentArrayFilterExpression; @@ -24,6 +28,7 @@ import org.hypertrace.core.documentstore.expression.impl.RelationalExpression; import org.hypertrace.core.documentstore.expression.operators.LogicalOperator; import org.hypertrace.core.documentstore.expression.type.FilterTypeExpression; +import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression; import org.hypertrace.core.documentstore.parser.FilterTypeExpressionVisitor; import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser; import org.hypertrace.core.documentstore.postgres.query.v1.parser.builder.PostgresSelectExpressionParserBuilder; @@ -117,7 +122,10 @@ WHERE TRIM('"' FROM elements::text) = 'Oxygen' switch (expression.getOperator()) { case ANY: return getFilterStringForAnyOperator(expression); - + case ALL: + return getFilterStringForAllOperator(expression); + case ONE: + return getFilterStringForOneOperator(expression); default: throw new UnsupportedOperationException( "Unsupported array operator: " + expression.getOperator()); @@ -137,7 +145,10 @@ FROM jsonb_array_elements(COALESCE(document->'planets', '[]'::jsonb)) AS planet switch (expression.getOperator()) { case ANY: return getFilterStringForAnyOperator(expression); - + case ALL: + return getFilterStringForAllOperator(expression); + case ONE: + return getFilterStringForOneOperator(expression); default: throw new UnsupportedOperationException( "Unsupported array operator: " + expression.getOperator()); @@ -332,4 +343,157 @@ private String getFilterStringForAnyOperator(final DocumentArrayFilterExpression parsedLhs, parsedLhs, alias, parsedFilter); } } + + /* + Native array (flat collection): + COALESCE(tags, ARRAY[]::text[]) @> ? -- bound as a typed array param, e.g. ['Blue','Green'] + + JSONB array (nested collection or JSONB column in a flat collection): + (CASE WHEN jsonb_typeof(colors) = 'array' THEN colors ELSE '[]'::jsonb END) @> ?::jsonb + -- bound as the JSON text '["Blue","Green"]' + */ + private String getFilterStringForAllOperator(final ArrayFilterExpression expression) { + final List values = getArrayOperatorFilterValues(expression); + final ArrayFieldContext fieldContext = getArrayFieldContext(expression); + + if (fieldContext.isNativeArray()) { + final PostgresDataType dataType = resolvePostgresDataType(values); + postgresQueryParser.getParamsBuilder().addArrayParam(values.toArray(), dataType.getSqlType()); + return String.format( + "COALESCE(%s, ARRAY[]%s) @> ?", fieldContext.parsedLhs(), dataType.getArrayTypeCast()); + } + + final String coalescedArray = + String.format( + "(CASE WHEN jsonb_typeof(%s) = 'array' THEN %s ELSE '[]'::jsonb END)", + fieldContext.parsedLhs(), fieldContext.parsedLhs()); + postgresQueryParser.getParamsBuilder().addObjectParam(toJsonArrayString(values)); + return String.format("%s @> ?::jsonb", coalescedArray); + } + + /* + Native array (flat collection): + array_length(COALESCE(tags, ARRAY[]::text[]), 1) = 1 + AND COALESCE(tags, ARRAY[]::text[]) && ? -- bound as a typed array param + + JSONB array (nested collection or JSONB column in a flat collection): + jsonb_array_length() = 1 + AND ( @> ?::jsonb OR @> ?::jsonb ...) + -- one bound single-element JSON array per filter value + */ + private String getFilterStringForOneOperator(final ArrayFilterExpression expression) { + final List values = getArrayOperatorFilterValues(expression); + final ArrayFieldContext fieldContext = getArrayFieldContext(expression); + + if (fieldContext.isNativeArray()) { + final PostgresDataType dataType = resolvePostgresDataType(values); + final String arrayTypeCast = dataType.getArrayTypeCast(); + final String coalescedArray = + String.format("COALESCE(%s, ARRAY[]%s)", fieldContext.parsedLhs(), arrayTypeCast); + postgresQueryParser.getParamsBuilder().addArrayParam(values.toArray(), dataType.getSqlType()); + return String.format("array_length(%s, 1) = 1 AND %s && ?", coalescedArray, coalescedArray); + } + + final String coalescedArray = + String.format( + "(CASE WHEN jsonb_typeof(%s) = 'array' THEN %s ELSE '[]'::jsonb END)", + fieldContext.parsedLhs(), fieldContext.parsedLhs()); + final String matchesAnyValue = + values.isEmpty() + ? "FALSE" + : values.stream() + .map( + value -> { + postgresQueryParser + .getParamsBuilder() + .addObjectParam(toJsonArrayString(List.of(value))); + return String.format("%s @> ?::jsonb", coalescedArray); + }) + .collect(Collectors.joining(" OR ")); + return String.format("jsonb_array_length(%s) = 1 AND (%s)", coalescedArray, matchesAnyValue); + } + + private ArrayFieldContext getArrayFieldContext(final ArrayFilterExpression expression) { + final boolean isFlatCollection = + postgresQueryParser.getPgColTransformer().getDocumentType() == DocumentType.FLAT; + final boolean isJsonbArray = expression.getArraySource() instanceof JsonIdentifierExpression; + final boolean isNativeArray = isFlatCollection && !isJsonbArray; + + final String identifierName = + expression + .getArraySource() + .accept(new PostgresIdentifierExpressionVisitor(postgresQueryParser)); + + final String parsedLhs; + if (isNativeArray) { + parsedLhs = postgresQueryParser.transformField(identifierName).getPgColumn(); + } else { + final PostgresIdentifierExpressionVisitor identifierVisitor = + new PostgresIdentifierExpressionVisitor(postgresQueryParser); + final PostgresSelectTypeExpressionVisitor arrayPathVisitor = + wrappingVisitorProvider == null + ? new PostgresFieldIdentifierExpressionVisitor(identifierVisitor) + : wrappingVisitorProvider.getForNonRelational(identifierVisitor); + parsedLhs = expression.getArraySource().accept(arrayPathVisitor); + } + + return new ArrayFieldContext(parsedLhs, isNativeArray); + } + + private List getArrayOperatorFilterValues(final ArrayFilterExpression expression) { + final FilterTypeExpression filter = expression.getFilter(); + if (!(filter instanceof RelationalExpression)) { + throw new UnsupportedOperationException( + "Array operator " + + expression.getOperator() + + " only supports a relational filter with a constant list of values, got: " + + filter); + } + + final SelectTypeExpression rhs = ((RelationalExpression) filter).getRhs(); + if (!(rhs instanceof ConstantExpression)) { + throw new UnsupportedOperationException( + "Array operator " + + expression.getOperator() + + " requires a constant list of values, got: " + + rhs); + } + + final Object value = ((ConstantExpression) rhs).getValue(); + return value instanceof List ? (List) value : List.of(value); + } + + private PostgresDataType resolvePostgresDataType(final List values) { + return values.stream() + .map(PostgresDataType::fromJavaValue) + .filter(dataType -> dataType != PostgresDataType.UNKNOWN) + .findFirst() + .orElse(PostgresDataType.TEXT); + } + + private String toJsonArrayString(final List values) { + try { + return new ObjectMapper().writeValueAsString(values); + } catch (JsonProcessingException e) { + throw new IllegalArgumentException("Unable to serialize array filter values: " + values, e); + } + } + + private static class ArrayFieldContext { + private final String parsedLhs; + private final boolean isNativeArray; + + ArrayFieldContext(final String parsedLhs, final boolean isNativeArray) { + this.parsedLhs = parsedLhs; + this.isNativeArray = isNativeArray; + } + + String parsedLhs() { + return parsedLhs; + } + + boolean isNativeArray() { + return isNativeArray; + } + } } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoArrayFilterParserTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoArrayFilterParserTest.java new file mode 100644 index 000000000..f7de01644 --- /dev/null +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/query/parser/MongoArrayFilterParserTest.java @@ -0,0 +1,179 @@ +package org.hypertrace.core.documentstore.mongo.query.parser; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import java.util.Map; +import org.hypertrace.core.documentstore.expression.impl.ArrayRelationalFilterExpression; +import org.hypertrace.core.documentstore.expression.impl.ConstantExpression; +import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression; +import org.hypertrace.core.documentstore.expression.impl.RelationalExpression; +import org.hypertrace.core.documentstore.expression.operators.ArrayOperator; +import org.hypertrace.core.documentstore.expression.operators.RelationalOperator; +import org.junit.jupiter.api.Test; + +class MongoArrayFilterParserTest { + + private final MongoFilterTypeExpressionParser parser = new MongoFilterTypeExpressionParser(); + + @Test + void testAllOperator() { + final ArrayRelationalFilterExpression expression = + ArrayRelationalFilterExpression.builder() + .operator(ArrayOperator.ALL) + .filter( + RelationalExpression.of( + IdentifierExpression.of("tags"), + RelationalOperator.IN, + ConstantExpression.ofStrings(List.of("Blue", "Green")))) + .build(); + + final Map result = parser.visit(expression); + + // {"$expr": {"$setIsSubset": [["Blue", "Green"], {"$ifNull": ["$tags", []]}]}} + final Map expr = getMap(result, "$expr"); + final List setIsSubset = getList(expr, "$setIsSubset"); + assertEquals(2, setIsSubset.size()); + assertEquals(List.of("Blue", "Green"), setIsSubset.get(0)); + + final Map ifNull = castToMap(setIsSubset.get(1)); + final Object[] ifNullArgs = (Object[]) ifNull.get("$ifNull"); + assertEquals("$tags", ifNullArgs[0]); + assertEquals(0, ((Object[]) ifNullArgs[1]).length); + } + + @Test + void testOneOperator() { + final ArrayRelationalFilterExpression expression = + ArrayRelationalFilterExpression.builder() + .operator(ArrayOperator.ONE) + .filter( + RelationalExpression.of( + IdentifierExpression.of("tags"), + RelationalOperator.IN, + ConstantExpression.ofStrings(List.of("Blue", "Green")))) + .build(); + + final Map result = parser.visit(expression); + + /* + {"$expr": {"$and": [ + {"$eq": [{"$size": {"$ifNull": ["$tags", []]}}, 1]}, + {"$in": [{"$arrayElemAt": [{"$ifNull": ["$tags", []]}, 0]}, ["Blue", "Green"]]} + ]}} + */ + final Map expr = getMap(result, "$expr"); + final List and = getList(expr, "$and"); + assertEquals(2, and.size()); + + final List eq = getList(castToMap(and.get(0)), "$eq"); + final Map size = castToMap(eq.get(0)); + final Map sizeIfNull = castToMap(size.get("$size")); + final Object[] sizeIfNullArgs = (Object[]) sizeIfNull.get("$ifNull"); + assertEquals("$tags", sizeIfNullArgs[0]); + assertEquals(1, eq.get(1)); + + final List in = getList(castToMap(and.get(1)), "$in"); + final Map arrayElemAt = castToMap(in.get(0)); + final List arrayElemAtArgs = getList(arrayElemAt, "$arrayElemAt"); + final Map elemIfNull = castToMap(arrayElemAtArgs.get(0)); + final Object[] elemIfNullArgs = (Object[]) elemIfNull.get("$ifNull"); + assertEquals("$tags", elemIfNullArgs[0]); + assertEquals(0, arrayElemAtArgs.get(1)); + assertEquals(List.of("Blue", "Green"), in.get(1)); + } + + @Test + void testAllOperatorWithSingleValue() { + final ArrayRelationalFilterExpression expression = + ArrayRelationalFilterExpression.builder() + .operator(ArrayOperator.ALL) + .filter( + RelationalExpression.of( + IdentifierExpression.of("tags"), + RelationalOperator.EQ, + ConstantExpression.of("Blue"))) + .build(); + + final Map result = parser.visit(expression); + + final Map expr = getMap(result, "$expr"); + final List setIsSubset = getList(expr, "$setIsSubset"); + assertEquals(List.of("Blue"), setIsSubset.get(0)); + } + + @Test + void testAnyOperatorStillSupported() { + final ArrayRelationalFilterExpression expression = + ArrayRelationalFilterExpression.builder() + .operator(ArrayOperator.ANY) + .filter( + RelationalExpression.of( + IdentifierExpression.of("tags"), + RelationalOperator.IN, + ConstantExpression.ofStrings(List.of("Blue", "Green")))) + .build(); + + // ANY remains supported + final Map result = parser.visit(expression); + assertTrue(result.containsKey("$expr")); + } + + @Test + void testAllOperatorInsideExprLocation() { + final MongoFilterTypeExpressionParser insideExprParser = + new MongoFilterTypeExpressionParser( + org.hypertrace.core.documentstore.mongo.query.parser.filter + .MongoRelationalFilterParserFactory.MongoRelationalFilterContext.builder() + .location( + org.hypertrace.core.documentstore.mongo.query.parser.filter + .MongoRelationalFilterParserFactory.FilterLocation.INSIDE_EXPR) + .build()); + + final ArrayRelationalFilterExpression expression = + ArrayRelationalFilterExpression.builder() + .operator(ArrayOperator.ALL) + .filter( + RelationalExpression.of( + IdentifierExpression.of("tags"), + RelationalOperator.IN, + ConstantExpression.ofStrings(List.of("Blue")))) + .build(); + + final Map result = insideExprParser.visit(expression); + + // already inside $expr, so no additional wrapping + assertTrue(result.containsKey("$setIsSubset")); + } + + @Test + void testOneOperatorRejectsNonConstantRhs() { + final ArrayRelationalFilterExpression expression = + ArrayRelationalFilterExpression.builder() + .operator(ArrayOperator.ONE) + .filter( + RelationalExpression.of( + IdentifierExpression.of("tags"), + RelationalOperator.EQ, + IdentifierExpression.of("otherField"))) + .build(); + + assertThrows(UnsupportedOperationException.class, () -> parser.visit(expression)); + } + + @SuppressWarnings("unchecked") + private Map castToMap(final Object object) { + return (Map) object; + } + + private Map getMap(final Map map, final String key) { + return castToMap(map.get(key)); + } + + @SuppressWarnings("unchecked") + private List getList(final Map map, final String key) { + return (List) map.get(key); + } +} 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..75d36d922 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 @@ -34,6 +34,7 @@ import org.hypertrace.core.documentstore.SingleValueKey; import org.hypertrace.core.documentstore.expression.impl.AggregateExpression; import org.hypertrace.core.documentstore.expression.impl.ArrayIdentifierExpression; +import org.hypertrace.core.documentstore.expression.impl.ArrayRelationalFilterExpression; import org.hypertrace.core.documentstore.expression.impl.ConstantExpression; import org.hypertrace.core.documentstore.expression.impl.FunctionExpression; import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression; @@ -43,6 +44,7 @@ import org.hypertrace.core.documentstore.expression.impl.LogicalExpression; import org.hypertrace.core.documentstore.expression.impl.RelationalExpression; import org.hypertrace.core.documentstore.expression.impl.UnnestExpression; +import org.hypertrace.core.documentstore.expression.operators.ArrayOperator; import org.hypertrace.core.documentstore.expression.operators.FunctionOperator; import org.hypertrace.core.documentstore.postgres.Params; import org.hypertrace.core.documentstore.postgres.PostgresTableIdentifier; @@ -1945,4 +1947,166 @@ void testNotInOperatorWithScalarStringField() { assertEquals("text", arrayParam.getSqlType()); assertEquals(2, arrayParam.getValues().length); } + + @Test + void testAllOperatorWithJsonbArrayField() { + Query query = + Query.builder() + .setFilter( + ArrayRelationalFilterExpression.builder() + .operator(ArrayOperator.ALL) + .filter( + RelationalExpression.of( + IdentifierExpression.of("tags"), + IN, + ConstantExpression.ofStrings(List.of("premium", "sale")))) + .build()) + .build(); + + PostgresQueryParser postgresQueryParser = + new PostgresQueryParser(TEST_TABLE, PostgresQueryTransformer.transform(query)); + + String sql = postgresQueryParser.parse(); + assertEquals( + "SELECT * FROM \"testCollection\" " + + "WHERE (CASE WHEN jsonb_typeof(document->'tags') = 'array' " + + "THEN document->'tags' ELSE '[]'::jsonb END) @> ?::jsonb", + sql); + + Params params = postgresQueryParser.getParamsBuilder().build(); + assertEquals(1, params.getObjectParams().size()); + assertEquals("[\"premium\",\"sale\"]", params.getObjectParams().get(1)); + } + + @Test + void testOneOperatorWithJsonbArrayField() { + Query query = + Query.builder() + .setFilter( + ArrayRelationalFilterExpression.builder() + .operator(ArrayOperator.ONE) + .filter( + RelationalExpression.of( + IdentifierExpression.of("tags"), + IN, + ConstantExpression.ofStrings(List.of("premium", "sale")))) + .build()) + .build(); + + PostgresQueryParser postgresQueryParser = + new PostgresQueryParser(TEST_TABLE, PostgresQueryTransformer.transform(query)); + + String sql = postgresQueryParser.parse(); + assertEquals( + "SELECT * FROM \"testCollection\" " + + "WHERE jsonb_array_length((CASE WHEN jsonb_typeof(document->'tags') = 'array' " + + "THEN document->'tags' ELSE '[]'::jsonb END)) = 1 " + + "AND ((CASE WHEN jsonb_typeof(document->'tags') = 'array' " + + "THEN document->'tags' ELSE '[]'::jsonb END) @> ?::jsonb " + + "OR (CASE WHEN jsonb_typeof(document->'tags') = 'array' " + + "THEN document->'tags' ELSE '[]'::jsonb END) @> ?::jsonb)", + sql); + + Params params = postgresQueryParser.getParamsBuilder().build(); + assertEquals(2, params.getObjectParams().size()); + assertEquals("[\"premium\"]", params.getObjectParams().get(1)); + assertEquals("[\"sale\"]", params.getObjectParams().get(2)); + } + + @Test + void testAllOperatorWithNativeArrayField() { + Query query = + Query.builder() + .setFilter( + ArrayRelationalFilterExpression.builder() + .operator(ArrayOperator.ALL) + .filter( + RelationalExpression.of( + ArrayIdentifierExpression.ofStrings("tags"), + IN, + ConstantExpression.ofStrings(List.of("premium", "sale")))) + .build()) + .build(); + + PostgresQueryParser postgresQueryParser = + new PostgresQueryParser( + TEST_TABLE, + PostgresQueryTransformer.transform(query), + new FlatPostgresFieldTransformer()); + + String sql = postgresQueryParser.parse(); + assertEquals( + "SELECT * FROM \"testCollection\" WHERE COALESCE(\"tags\", ARRAY[]::text[]) @> ?", sql); + + Params params = postgresQueryParser.getParamsBuilder().build(); + assertEquals(1, params.getObjectParams().size()); + Params.ArrayParam arrayParam = (Params.ArrayParam) params.getObjectParams().get(1); + assertEquals("text", arrayParam.getSqlType()); + assertEquals(2, arrayParam.getValues().length); + } + + @Test + void testOneOperatorWithNativeArrayField() { + Query query = + Query.builder() + .setFilter( + ArrayRelationalFilterExpression.builder() + .operator(ArrayOperator.ONE) + .filter( + RelationalExpression.of( + ArrayIdentifierExpression.ofStrings("tags"), + IN, + ConstantExpression.ofStrings(List.of("premium", "sale")))) + .build()) + .build(); + + PostgresQueryParser postgresQueryParser = + new PostgresQueryParser( + TEST_TABLE, + PostgresQueryTransformer.transform(query), + new FlatPostgresFieldTransformer()); + + String sql = postgresQueryParser.parse(); + assertEquals( + "SELECT * FROM \"testCollection\" " + + "WHERE array_length(COALESCE(\"tags\", ARRAY[]::text[]), 1) = 1 " + + "AND COALESCE(\"tags\", ARRAY[]::text[]) && ?", + sql); + + Params params = postgresQueryParser.getParamsBuilder().build(); + assertEquals(1, params.getObjectParams().size()); + Params.ArrayParam arrayParam = (Params.ArrayParam) params.getObjectParams().get(1); + assertEquals("text", arrayParam.getSqlType()); + assertEquals(2, arrayParam.getValues().length); + } + + @Test + void testAllOperatorWithNestedJsonbArrayField() { + Query query = + Query.builder() + .setFilter( + ArrayRelationalFilterExpression.builder() + .operator(ArrayOperator.ALL) + .filter( + RelationalExpression.of( + IdentifierExpression.of("scope.environmentScope.environmentIds"), + IN, + ConstantExpression.ofStrings(List.of("env-1", "env-2")))) + .build()) + .build(); + + PostgresQueryParser postgresQueryParser = + new PostgresQueryParser(TEST_TABLE, PostgresQueryTransformer.transform(query)); + + String sql = postgresQueryParser.parse(); + assertEquals( + "SELECT * FROM \"testCollection\" " + + "WHERE (CASE WHEN jsonb_typeof(document->'scope'->'environmentScope'->'environmentIds') = 'array' " + + "THEN document->'scope'->'environmentScope'->'environmentIds' ELSE '[]'::jsonb END) @> ?::jsonb", + sql); + + Params params = postgresQueryParser.getParamsBuilder().build(); + assertEquals(1, params.getObjectParams().size()); + assertEquals("[\"env-1\",\"env-2\"]", params.getObjectParams().get(1)); + } }