Skip to content

Commit 4192bde

Browse files
suddendustsuddendust
authored
Relational operators on nested fields in flat collections (#243)
1 parent 704c275 commit 4192bde

File tree

5 files changed

+383
-17
lines changed

5 files changed

+383
-17
lines changed

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

Lines changed: 214 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4046,6 +4046,220 @@ void testFlatCollectionArrayAnyOnJsonbArray(String dataStoreName) {
40464046
// ids 1 and 5 have "Blue" in their colors array
40474047
assertEquals(2, count, "Should find 2 items with 'Blue' color (ids 1, 5)");
40484048
}
4049+
4050+
/**
4051+
* Tests for relational operators on JSONB nested fields in flat collections. Tests: CONTAINS,
4052+
* NOT_CONTAINS, IN, NOT_IN, EQ, NEQ, LT, GT on JSONB columns.
4053+
*/
4054+
@Nested
4055+
class FlatCollectionJsonbRelationalOperatorTest {
4056+
4057+
/**
4058+
* Tests CONTAINS and NOT_CONTAINS operators on JSONB array fields. - CONTAINS: finds
4059+
* documents where array contains the value - NOT_CONTAINS: finds documents where array
4060+
* doesn't contain the value (including NULL)
4061+
*/
4062+
@ParameterizedTest
4063+
@ArgumentsSource(PostgresProvider.class)
4064+
void testJsonbArrayContainsOperators(String dataStoreName) {
4065+
Datastore datastore = datastoreMap.get(dataStoreName);
4066+
Collection flatCollection =
4067+
datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);
4068+
4069+
// Test 1: CONTAINS - props.colors CONTAINS "Green"
4070+
// Expected: 1 document (id=1, Dettol Soap has ["Green", "White"])
4071+
Query containsQuery =
4072+
Query.builder()
4073+
.setFilter(
4074+
RelationalExpression.of(
4075+
JsonIdentifierExpression.of("props", "colors"),
4076+
CONTAINS,
4077+
ConstantExpression.of("Green")))
4078+
.build();
4079+
4080+
long containsCount = flatCollection.count(containsQuery);
4081+
assertEquals(1, containsCount, "CONTAINS: Should find 1 document with Green color");
4082+
4083+
// Test 2: NOT_CONTAINS - props.colors NOT_CONTAINS "Green" AND _id <= 8
4084+
// Expected: 7 documents (all except id=1 which has Green, limited to first 8)
4085+
Query notContainsQuery =
4086+
Query.builder()
4087+
.setFilter(
4088+
LogicalExpression.builder()
4089+
.operator(LogicalOperator.AND)
4090+
.operand(
4091+
RelationalExpression.of(
4092+
JsonIdentifierExpression.of("props", "colors"),
4093+
NOT_CONTAINS,
4094+
ConstantExpression.of("Green")))
4095+
.operand(
4096+
RelationalExpression.of(
4097+
IdentifierExpression.of("_id"), LTE, ConstantExpression.of(8)))
4098+
.build())
4099+
.build();
4100+
4101+
long notContainsCount = flatCollection.count(notContainsQuery);
4102+
assertEquals(
4103+
7, notContainsCount, "NOT_CONTAINS: Should find 7 documents without Green color");
4104+
}
4105+
4106+
/**
4107+
* Tests IN and NOT_IN operators on JSONB scalar fields. - IN: finds documents where field
4108+
* value is in the provided list - NOT_IN: finds documents where field value is not in the
4109+
* list (including NULL)
4110+
*/
4111+
@ParameterizedTest
4112+
@ArgumentsSource(PostgresProvider.class)
4113+
void testJsonbScalarInOperators(String dataStoreName) {
4114+
Datastore datastore = datastoreMap.get(dataStoreName);
4115+
Collection flatCollection =
4116+
datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);
4117+
4118+
// Test 1: IN - props.brand IN ["Dettol", "Lifebuoy"]
4119+
// Expected: 2 documents (id=1 Dettol, id=5 Lifebuoy)
4120+
Query inQuery =
4121+
Query.builder()
4122+
.setFilter(
4123+
RelationalExpression.of(
4124+
JsonIdentifierExpression.of("props", "brand"),
4125+
IN,
4126+
ConstantExpression.ofStrings(List.of("Dettol", "Lifebuoy"))))
4127+
.build();
4128+
4129+
long inCount = flatCollection.count(inQuery);
4130+
assertEquals(2, inCount, "IN: Should find 2 documents with Dettol or Lifebuoy brand");
4131+
4132+
// Test 2: NOT_IN - props.brand NOT_IN ["Dettol"] AND _id <= 8
4133+
// Expected: 7 documents (all except id=1 which is Dettol, limited to first 8)
4134+
Query notInQuery =
4135+
Query.builder()
4136+
.setFilter(
4137+
LogicalExpression.builder()
4138+
.operator(LogicalOperator.AND)
4139+
.operand(
4140+
RelationalExpression.of(
4141+
JsonIdentifierExpression.of("props", "brand"),
4142+
NOT_IN,
4143+
ConstantExpression.ofStrings(List.of("Dettol"))))
4144+
.operand(
4145+
RelationalExpression.of(
4146+
IdentifierExpression.of("_id"), LTE, ConstantExpression.of(8)))
4147+
.build())
4148+
.build();
4149+
4150+
long notInCount = flatCollection.count(notInQuery);
4151+
assertEquals(7, notInCount, "NOT_IN: Should find 7 documents without Dettol brand");
4152+
}
4153+
4154+
/**
4155+
* Tests EQ and NEQ operators on JSONB scalar fields. - EQ: finds documents where field equals
4156+
* the value - NEQ: finds documents where field doesn't equal the value (excluding NULL)
4157+
*/
4158+
@ParameterizedTest
4159+
@ArgumentsSource(PostgresProvider.class)
4160+
void testJsonbScalarEqualityOperators(String dataStoreName) {
4161+
Datastore datastore = datastoreMap.get(dataStoreName);
4162+
Collection flatCollection =
4163+
datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);
4164+
4165+
// Test 1: EQ - props.brand EQ "Dettol"
4166+
// Expected: 1 document (id=1, Dettol Soap)
4167+
Query eqQuery =
4168+
Query.builder()
4169+
.setFilter(
4170+
RelationalExpression.of(
4171+
JsonIdentifierExpression.of("props", "brand"),
4172+
EQ,
4173+
ConstantExpression.of("Dettol")))
4174+
.build();
4175+
4176+
long eqCount = flatCollection.count(eqQuery);
4177+
assertEquals(1, eqCount, "EQ: Should find 1 document with Dettol brand");
4178+
4179+
// Test 2: NEQ - props.brand NEQ "Dettol" (no _id filter needed)
4180+
// Expected: 2 documents (id=3 Sunsilk, id=5 Lifebuoy, excluding NULL props)
4181+
Query neqQuery =
4182+
Query.builder()
4183+
.setFilter(
4184+
RelationalExpression.of(
4185+
JsonIdentifierExpression.of("props", "brand"),
4186+
NEQ,
4187+
ConstantExpression.of("Dettol")))
4188+
.build();
4189+
4190+
long neqCount = flatCollection.count(neqQuery);
4191+
assertEquals(2, neqCount, "NEQ: Should find 2 documents without Dettol brand");
4192+
}
4193+
4194+
/**
4195+
* Tests LT, GT, LTE, GTE comparison operators on JSONB numeric fields. Tests deeply nested
4196+
* numeric fields like props.seller.address.pincode. Data: ids 1,3 have pincode 400004; ids
4197+
* 5,7 have pincode 700007; rest are NULL
4198+
*/
4199+
@ParameterizedTest
4200+
@ArgumentsSource(PostgresProvider.class)
4201+
void testJsonbNumericComparisonOperators(String dataStoreName) {
4202+
Datastore datastore = datastoreMap.get(dataStoreName);
4203+
Collection flatCollection =
4204+
datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);
4205+
4206+
// Test 1: GT - props.seller.address.pincode > 500000
4207+
// Expected: 2 documents (ids 5,7 with pincode 700007 in Kolkata)
4208+
Query gtQuery =
4209+
Query.builder()
4210+
.setFilter(
4211+
RelationalExpression.of(
4212+
JsonIdentifierExpression.of("props", "seller", "address", "pincode"),
4213+
GT,
4214+
ConstantExpression.of(500000)))
4215+
.build();
4216+
4217+
long gtCount = flatCollection.count(gtQuery);
4218+
assertEquals(2, gtCount, "GT: Should find 2 documents with pincode > 500000");
4219+
4220+
// Test 2: LT - props.seller.address.pincode < 500000
4221+
// Expected: 2 documents (ids 1,3 with pincode 400004 in Mumbai)
4222+
Query ltQuery =
4223+
Query.builder()
4224+
.setFilter(
4225+
RelationalExpression.of(
4226+
JsonIdentifierExpression.of("props", "seller", "address", "pincode"),
4227+
LT,
4228+
ConstantExpression.of(500000)))
4229+
.build();
4230+
4231+
long ltCount = flatCollection.count(ltQuery);
4232+
assertEquals(2, ltCount, "LT: Should find 2 documents with pincode < 500000");
4233+
4234+
// Test 3: GTE - props.seller.address.pincode >= 700000
4235+
// Expected: 2 documents (ids 5,7 with pincode 700007)
4236+
Query gteQuery =
4237+
Query.builder()
4238+
.setFilter(
4239+
RelationalExpression.of(
4240+
JsonIdentifierExpression.of("props", "seller", "address", "pincode"),
4241+
GTE,
4242+
ConstantExpression.of(700000)))
4243+
.build();
4244+
4245+
long gteCount = flatCollection.count(gteQuery);
4246+
assertEquals(2, gteCount, "GTE: Should find 2 documents with pincode >= 700000");
4247+
4248+
// Test 4: LTE - props.seller.address.pincode <= 400004
4249+
// Expected: 2 documents (ids 1,3 with pincode 400004)
4250+
Query lteQuery =
4251+
Query.builder()
4252+
.setFilter(
4253+
RelationalExpression.of(
4254+
JsonIdentifierExpression.of("props", "seller", "address", "pincode"),
4255+
LTE,
4256+
ConstantExpression.of(400004)))
4257+
.build();
4258+
4259+
long lteCount = flatCollection.count(lteQuery);
4260+
assertEquals(2, lteCount, "LTE: Should find 2 documents with pincode <= 400004");
4261+
}
4262+
}
40494263
}
40504264

40514265
@Nested

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotContainsRelationalFilterParser.java‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.hypertrace.core.documentstore.postgres.query.v1.parser.filter;
22

33
import org.hypertrace.core.documentstore.DocumentType;
4+
import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
45
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
56
import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresContainsRelationalFilterParserNonJsonField;
67

@@ -16,16 +17,26 @@ public String parse(
1617
final RelationalExpression expression, final PostgresRelationalFilterContext context) {
1718
final String parsedLhs = expression.getLhs().accept(context.lhsParser());
1819

19-
boolean isFirstClassField =
20-
context.getPgColTransformer().getDocumentType() == DocumentType.FLAT;
21-
if (isFirstClassField) {
20+
boolean useJsonParser = shouldUseJsonParser(expression, context);
21+
22+
if (useJsonParser) {
23+
// Use the JSON logic for JSON document fields
24+
jsonContainsParser.parse(expression, context); // This adds the parameter.
25+
return String.format("%s IS NULL OR NOT %s @> ?::jsonb", parsedLhs, parsedLhs);
26+
} else {
2227
// Use the non-JSON logic for first-class fields
2328
String containsExpression = nonJsonContainsParser.parse(expression, context);
2429
return String.format("%s IS NULL OR NOT (%s)", parsedLhs, containsExpression);
25-
} else {
26-
// Use the JSON logic for document fields.
27-
jsonContainsParser.parse(expression, context); // This adds the parameter.
28-
return String.format("%s IS NULL OR NOT %s @> ?::jsonb", parsedLhs, parsedLhs);
2930
}
3031
}
32+
33+
private boolean shouldUseJsonParser(
34+
final RelationalExpression expression, final PostgresRelationalFilterContext context) {
35+
36+
boolean isJsonField = expression.getLhs() instanceof JsonIdentifierExpression;
37+
boolean isFlatCollection = context.getPgColTransformer().getDocumentType() == DocumentType.FLAT;
38+
boolean useJsonParser = !isFlatCollection || isJsonField;
39+
40+
return useJsonParser;
41+
}
3142
}

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresNotInRelationalFilterParser.java‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.hypertrace.core.documentstore.postgres.query.v1.parser.filter;
22

33
import org.hypertrace.core.documentstore.DocumentType;
4+
import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
45
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
56
import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresInRelationalFilterParserNonJsonField;
67

@@ -16,17 +17,26 @@ public String parse(
1617
final RelationalExpression expression, final PostgresRelationalFilterContext context) {
1718
final String parsedLhs = expression.getLhs().accept(context.lhsParser());
1819

19-
PostgresInRelationalFilterParserInterface inFilterParser = getInFilterParser(context);
20+
PostgresInRelationalFilterParserInterface inFilterParser =
21+
getInFilterParser(expression, context);
2022

2123
final String parsedInExpression = inFilterParser.parse(expression, context);
2224
return String.format("%s IS NULL OR NOT (%s)", parsedLhs, parsedInExpression);
2325
}
2426

2527
private PostgresInRelationalFilterParserInterface getInFilterParser(
26-
PostgresRelationalFilterContext context) {
27-
boolean isFirstClassField =
28-
context.getPgColTransformer().getDocumentType() == DocumentType.FLAT;
28+
final RelationalExpression expression, PostgresRelationalFilterContext context) {
29+
// Check if LHS is a JSON field (JSONB column access)
30+
boolean isJsonField = expression.getLhs() instanceof JsonIdentifierExpression;
2931

30-
return isFirstClassField ? nonJsonFieldInFilterParser : jsonFieldInFilterParser;
32+
// Check if the collection type is flat or nested
33+
boolean isFlatCollection = context.getPgColTransformer().getDocumentType() == DocumentType.FLAT;
34+
35+
// Use JSON parser for:
36+
// 1. Nested collections - !isFlatCollection
37+
// 2. JSON fields within flat collections - isJsonField
38+
boolean useJsonParser = !isFlatCollection || isJsonField;
39+
40+
return useJsonParser ? jsonFieldInFilterParser : nonJsonFieldInFilterParser;
3141
}
3242
}

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresRelationalFilterParserFactoryImpl.java‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,13 @@
1212

1313
import com.google.common.collect.Maps;
1414
import java.util.Map;
15+
import org.hypertrace.core.documentstore.DocumentType;
16+
import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
1517
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
1618
import org.hypertrace.core.documentstore.expression.operators.RelationalOperator;
1719
import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser;
1820
import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresContainsRelationalFilterParserNonJsonField;
1921
import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresInRelationalFilterParserNonJsonField;
20-
import org.hypertrace.core.documentstore.postgres.query.v1.transformer.FlatPostgresFieldTransformer;
2122

2223
public class PostgresRelationalFilterParserFactoryImpl
2324
implements PostgresRelationalFilterParserFactory {
@@ -52,13 +53,19 @@ public class PostgresRelationalFilterParserFactoryImpl
5253
public PostgresRelationalFilterParser parser(
5354
final RelationalExpression expression, final PostgresQueryParser postgresQueryParser) {
5455

55-
boolean isFirstClassField =
56-
postgresQueryParser.getPgColTransformer() instanceof FlatPostgresFieldTransformer;
56+
// Check if LHS is a JSON field (JSONB column access)
57+
boolean isJsonField = expression.getLhs() instanceof JsonIdentifierExpression;
58+
59+
// Check if the collection type is flat or nested
60+
boolean isFlatCollection =
61+
postgresQueryParser.getPgColTransformer().getDocumentType() == DocumentType.FLAT;
62+
63+
boolean useJsonParser = !isFlatCollection || isJsonField;
5764

5865
if (expression.getOperator() == CONTAINS) {
59-
return isFirstClassField ? nonJsonFieldContainsParser : jsonFieldContainsParser;
66+
return useJsonParser ? jsonFieldContainsParser : nonJsonFieldContainsParser;
6067
} else if (expression.getOperator() == IN) {
61-
return isFirstClassField ? nonJsonFieldInFilterParser : jsonFieldInFilterParser;
68+
return useJsonParser ? jsonFieldInFilterParser : nonJsonFieldInFilterParser;
6269
}
6370

6471
return parserMap.getOrDefault(expression.getOperator(), postgresStandardRelationalFilterParser);

0 commit comments

Comments
 (0)