Skip to content

Commit 079326d

Browse files
suddendustsuddendust
committed
Add more UTs
1 parent 2c70154 commit 079326d

File tree

5 files changed

+263
-28
lines changed

5 files changed

+263
-28
lines changed

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4828,9 +4828,10 @@ void testEqNotEqScalar(String dataStoreName) {
48284828
assertEquals(7, count);
48294829
}
48304830

4831+
/** Tests the behavior of EQ/NEQ on array fields with array RHS. */
48314832
@ParameterizedTest
48324833
@ArgumentsSource(PostgresProvider.class)
4833-
void testEqNotEqArrays(String dataStoreName) {
4834+
void testEqAndNeqOnArrays(String dataStoreName) {
48344835
Collection flatCollection = getFlatCollection(dataStoreName);
48354836

48364837
Query eqQuery =
@@ -5581,9 +5582,13 @@ void testContainsNotContainsArraysInArrays(String dataStoreName) {
55815582
assertEquals(9, count);
55825583
}
55835584

5585+
/**
5586+
* Given LHS is a nested array, operator is EQ/NEQ, and RHS is a scalar, it should behave as
5587+
* containment (LHS contains RHS)
5588+
*/
55845589
@ParameterizedTest
55855590
@ArgumentsSource(PostgresProvider.class)
5586-
void testEqNotEqScalarsInArrays(String dataStoreName) {
5591+
void testEqNeqScalarsOnArrays(String dataStoreName) {
55875592
Collection flatCollection = getFlatCollection(dataStoreName);
55885593

55895594
// Should be treated like CONTAINS

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

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import org.hypertrace.core.documentstore.expression.impl.JsonFieldType;
1010
import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
1111
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
12+
import org.hypertrace.core.documentstore.expression.type.SelectTypeExpression;
1213
import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser;
1314
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresConstantExpressionVisitor;
1415
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresDataAccessorIdentifierExpressionVisitor;
@@ -27,9 +28,6 @@ public PostgresSelectExpressionParserBuilderImpl(PostgresQueryParser postgresQue
2728

2829
@Override
2930
public PostgresSelectTypeExpressionVisitor build(final RelationalExpression expression) {
30-
// For EQ/NEQ on array fields, treat like CONTAINS to use -> instead of ->>
31-
boolean isEqOrNeqOnArrayField = isEqOrNeqOnArrayField(expression);
32-
3331
switch (expression.getOperator()) {
3432
case CONTAINS:
3533
case NOT_CONTAINS:
@@ -42,13 +40,13 @@ public PostgresSelectTypeExpressionVisitor build(final RelationalExpression expr
4240

4341
case EQ:
4442
case NEQ:
45-
if (isEqOrNeqOnArrayField) {
43+
// For EQ/NEQ on array fields, treat like CONTAINS to use -> instead of ->>
44+
if (shouldSwitchToContainsFlow(expression)) {
4645
// Use field identifier (JSON accessor ->) for array fields
4746
return new PostgresFunctionExpressionVisitor(
4847
new PostgresFieldIdentifierExpressionVisitor(this.postgresQueryParser));
4948
}
5049
// Fall through to default for non-array fields
51-
5250
default:
5351
return new PostgresFunctionExpressionVisitor(
5452
new PostgresDataAccessorIdentifierExpressionVisitor(
@@ -66,11 +64,11 @@ public PostgresSelectTypeExpressionVisitor build(final RelationalExpression expr
6664
* <p>Handles both:
6765
*
6866
* <ul>
69-
* <li>JsonIdentifierExpression with array field type (JSONB arrays)
70-
* <li>ArrayIdentifierExpression with array type (top-level array columns)
67+
* <li>{@link JsonIdentifierExpression} with array field type (JSONB arrays)
68+
* <li>{@link ArrayIdentifierExpression} with array type (top-level array columns)
7169
* </ul>
7270
*/
73-
private boolean isEqOrNeqOnArrayField(final RelationalExpression expression) {
71+
private boolean shouldSwitchToContainsFlow(final RelationalExpression expression) {
7472
if (expression.getOperator() != EQ && expression.getOperator() != NEQ) {
7573
return false;
7674
}
@@ -87,9 +85,7 @@ private boolean isEqOrNeqOnArrayField(final RelationalExpression expression) {
8785
return isArrayField(expression.getLhs());
8886
}
8987

90-
/** Checks if the expression is an array field. */
91-
private boolean isArrayField(
92-
final org.hypertrace.core.documentstore.expression.type.SelectTypeExpression lhs) {
88+
private boolean isArrayField(final SelectTypeExpression lhs) {
9389
if (lhs instanceof JsonIdentifierExpression) {
9490
JsonIdentifierExpression jsonExpr = (JsonIdentifierExpression) lhs;
9591
return jsonExpr

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

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,8 @@ public PostgresRelationalFilterParser parser(
6464
return parserMap.get(NOT_CONTAINS);
6565
}
6666

67-
// For EQ/NEQ on array fields with array RHS, use specialized array equality parser
67+
// For EQ/NEQ on array fields with array RHS, use specialized array equality parser (exact match
68+
// instead of containment)
6869
if (shouldUseArrayEqualityParser(expression, postgresQueryParser)) {
6970
return expression.getLhs().accept(new PostgresArrayEqualityParserSelector());
7071
}
@@ -110,11 +111,10 @@ private boolean shouldConvertEqToContains(
110111

111112
// Check if field has been unnested - unnested fields are scalar, not arrays
112113
String fieldName = getFieldName(expression.getLhs());
113-
if (fieldName != null && postgresQueryParser.getPgColumnNames().containsKey(fieldName)) {
114-
return false; // Field is unnested - treat as scalar
115-
}
116-
117-
return true;
114+
return fieldName == null
115+
|| !postgresQueryParser
116+
.getPgColumnNames()
117+
.containsKey(fieldName); // Field is unnested - treat as scalar
118118
}
119119

120120
/**
@@ -124,8 +124,9 @@ private boolean shouldConvertEqToContains(
124124
*
125125
* <ul>
126126
* <li>Operator is EQ or NEQ
127-
* <li>RHS is an array/iterable (for exact match)
128-
* <li>LHS is either JsonIdentifierExpression with array type OR ArrayIdentifierExpression
127+
* <li>RHS is an array/iterable (for exact match).
128+
* <li>LHS is either {@link JsonIdentifierExpression} with array type OR {@link
129+
* ArrayIdentifierExpression}
129130
* <li>Field has NOT been unnested (unnested fields are scalar, not arrays)
130131
* </ul>
131132
*/
@@ -142,14 +143,13 @@ private boolean shouldUseArrayEqualityParser(
142143

143144
// Check if field has been unnested - unnested fields are scalar, not arrays
144145
String fieldName = getFieldName(expression.getLhs());
145-
if (fieldName != null && postgresQueryParser.getPgColumnNames().containsKey(fieldName)) {
146-
return false;
147-
}
148-
149-
return true;
146+
return fieldName == null || !postgresQueryParser.getPgColumnNames().containsKey(fieldName);
150147
}
151148

152-
/** Checks if the RHS expression contains an array/iterable value. */
149+
/**
150+
* Checks if the RHS expression contains an array/iterable value. Currently, we don't have a very
151+
* clean way to get the RHS data type. //todo: Implement a clean way to get the RHS data type
152+
*/
153153
private boolean isArrayRhs(final SelectTypeExpression rhs) {
154154
if (rhs instanceof ConstantExpression) {
155155
ConstantExpression constExpr = (ConstantExpression) rhs;
@@ -172,7 +172,7 @@ private boolean isArrayField(final SelectTypeExpression lhs) {
172172
|| fieldType == JsonFieldType.OBJECT_ARRAY)
173173
.orElse(false);
174174
}
175-
return lhs instanceof ArrayIdentifierExpression;
175+
return !(lhs instanceof ArrayIdentifierExpression);
176176
}
177177

178178
/** Extracts the field name from an identifier expression. */

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

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
package org.hypertrace.core.documentstore.postgres.query.v1.parser.filter;
2+
3+
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ;
4+
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.NEQ;
5+
import static org.junit.jupiter.api.Assertions.assertEquals;
6+
import static org.mockito.ArgumentMatchers.any;
7+
import static org.mockito.Mockito.mock;
8+
import static org.mockito.Mockito.when;
9+
10+
import java.util.List;
11+
import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
12+
import org.hypertrace.core.documentstore.expression.impl.JsonFieldType;
13+
import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression;
14+
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
15+
import org.hypertrace.core.documentstore.postgres.Params;
16+
import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser;
17+
import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.PostgresRelationalFilterParser.PostgresRelationalFilterContext;
18+
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresSelectTypeExpressionVisitor;
19+
import org.junit.jupiter.api.BeforeEach;
20+
import org.junit.jupiter.api.Test;
21+
22+
class PostgresJsonArrayEqualityFilterParserTest {
23+
24+
private PostgresJsonArrayEqualityFilterParser parser;
25+
private PostgresRelationalFilterContext context;
26+
private PostgresSelectTypeExpressionVisitor lhsParser;
27+
private PostgresQueryParser queryParser;
28+
private Params.Builder paramsBuilder;
29+
30+
@BeforeEach
31+
void setUp() {
32+
parser = new PostgresJsonArrayEqualityFilterParser();
33+
lhsParser = mock(PostgresSelectTypeExpressionVisitor.class);
34+
queryParser = mock(PostgresQueryParser.class);
35+
paramsBuilder = Params.newBuilder();
36+
37+
when(queryParser.getParamsBuilder()).thenReturn(paramsBuilder);
38+
39+
context =
40+
PostgresRelationalFilterContext.builder()
41+
.lhsParser(lhsParser)
42+
.postgresQueryParser(queryParser)
43+
.build();
44+
}
45+
46+
@Test
47+
void testParse_nullRhs() {
48+
JsonIdentifierExpression lhs =
49+
JsonIdentifierExpression.of("props", JsonFieldType.STRING_ARRAY, "colors");
50+
// RHS parsed value will be null from rhsParser
51+
ConstantExpression rhs = ConstantExpression.of("ignored");
52+
RelationalExpression expression = RelationalExpression.of(lhs, EQ, rhs);
53+
54+
when(lhsParser.visit(any(JsonIdentifierExpression.class))).thenReturn("props->'colors'");
55+
// Simulate rhsParser returning null
56+
PostgresSelectTypeExpressionVisitor rhsParser = mock(PostgresSelectTypeExpressionVisitor.class);
57+
when(rhsParser.visit(rhs)).thenReturn(null);
58+
context =
59+
PostgresRelationalFilterContext.builder()
60+
.lhsParser(lhsParser)
61+
.rhsParser(rhsParser)
62+
.postgresQueryParser(queryParser)
63+
.build();
64+
65+
String result = parser.parse(expression, context);
66+
67+
assertEquals("props->'colors' = NULL", result);
68+
assertEquals(0, paramsBuilder.build().getObjectParams().size());
69+
}
70+
71+
@Test
72+
void testParseScalarRhsEq() {
73+
JsonIdentifierExpression lhs =
74+
JsonIdentifierExpression.of("props", JsonFieldType.STRING_ARRAY, "colors");
75+
ConstantExpression rhs = ConstantExpression.of("Blue");
76+
RelationalExpression expression = RelationalExpression.of(lhs, EQ, rhs);
77+
78+
when(lhsParser.visit(any(JsonIdentifierExpression.class))).thenReturn("props->'colors'");
79+
80+
String result = parser.parse(expression, context);
81+
82+
assertEquals("props->'colors' @> ?::jsonb", result);
83+
assertEquals(1, paramsBuilder.build().getObjectParams().size());
84+
assertEquals("Blue", paramsBuilder.build().getObjectParams().get(1));
85+
}
86+
87+
@Test
88+
void testParseIterableRhsEq() {
89+
JsonIdentifierExpression lhs =
90+
JsonIdentifierExpression.of("props", JsonFieldType.STRING_ARRAY, "colors");
91+
ConstantExpression rhs = ConstantExpression.ofStrings(List.of("Blue", "Green"));
92+
RelationalExpression expression = RelationalExpression.of(lhs, EQ, rhs);
93+
94+
when(lhsParser.visit(any(JsonIdentifierExpression.class))).thenReturn("props->'colors'");
95+
96+
String result = parser.parse(expression, context);
97+
98+
assertEquals("props->'colors' = ?::jsonb", result);
99+
assertEquals(1, paramsBuilder.build().getObjectParams().size());
100+
assertEquals("[\"Blue\",\"Green\"]", paramsBuilder.build().getObjectParams().get(1));
101+
}
102+
103+
/** Tests parsing when RHS is an iterable (e.g., list of strings) and the operator is NEQ */
104+
@Test
105+
void testParseIterableRhsNeq() {
106+
JsonIdentifierExpression lhs =
107+
JsonIdentifierExpression.of("props", JsonFieldType.STRING_ARRAY, "colors");
108+
ConstantExpression rhs = ConstantExpression.ofStrings(List.of("Blue", "Red"));
109+
RelationalExpression expression = RelationalExpression.of(lhs, NEQ, rhs);
110+
111+
when(lhsParser.visit(any(JsonIdentifierExpression.class))).thenReturn("props->'colors'");
112+
113+
String result = parser.parse(expression, context);
114+
115+
assertEquals("props->'colors' != ?::jsonb", result);
116+
assertEquals(1, paramsBuilder.build().getObjectParams().size());
117+
assertEquals("[\"Blue\",\"Red\"]", paramsBuilder.build().getObjectParams().get(1));
118+
}
119+
}

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

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
package org.hypertrace.core.documentstore.postgres.query.v1.parser.filter;
2+
3+
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ;
4+
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.NEQ;
5+
import static org.junit.jupiter.api.Assertions.assertEquals;
6+
import static org.mockito.ArgumentMatchers.any;
7+
import static org.mockito.Mockito.mock;
8+
import static org.mockito.Mockito.when;
9+
10+
import java.util.List;
11+
import org.hypertrace.core.documentstore.expression.impl.ArrayIdentifierExpression;
12+
import org.hypertrace.core.documentstore.expression.impl.ArrayType;
13+
import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
14+
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
15+
import org.hypertrace.core.documentstore.postgres.Params;
16+
import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser;
17+
import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.PostgresRelationalFilterParser.PostgresRelationalFilterContext;
18+
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresSelectTypeExpressionVisitor;
19+
import org.junit.jupiter.api.BeforeEach;
20+
import org.junit.jupiter.api.Test;
21+
22+
class PostgresTopLevelArrayEqualityFilterParserTest {
23+
24+
private PostgresTopLevelArrayEqualityFilterParser parser;
25+
private PostgresRelationalFilterContext context;
26+
private PostgresSelectTypeExpressionVisitor lhsParser;
27+
private PostgresQueryParser queryParser;
28+
private Params.Builder paramsBuilder;
29+
30+
@BeforeEach
31+
void setUp() {
32+
parser = new PostgresTopLevelArrayEqualityFilterParser();
33+
lhsParser = mock(PostgresSelectTypeExpressionVisitor.class);
34+
queryParser = mock(PostgresQueryParser.class);
35+
paramsBuilder = Params.newBuilder();
36+
37+
when(queryParser.getParamsBuilder()).thenReturn(paramsBuilder);
38+
39+
context =
40+
PostgresRelationalFilterContext.builder()
41+
.lhsParser(lhsParser)
42+
.postgresQueryParser(queryParser)
43+
.build();
44+
}
45+
46+
@Test
47+
void testParseNullRhs() {
48+
ArrayIdentifierExpression lhs = ArrayIdentifierExpression.of("tags", ArrayType.TEXT);
49+
ConstantExpression rhs = ConstantExpression.of("ignored");
50+
RelationalExpression expression = RelationalExpression.of(lhs, EQ, rhs);
51+
52+
when(lhsParser.visit(any(ArrayIdentifierExpression.class))).thenReturn("tags");
53+
// Simulate rhsParser returning null
54+
PostgresSelectTypeExpressionVisitor rhsParser = mock(PostgresSelectTypeExpressionVisitor.class);
55+
when(rhsParser.visit(rhs)).thenReturn(null);
56+
context =
57+
PostgresRelationalFilterContext.builder()
58+
.lhsParser(lhsParser)
59+
.rhsParser(rhsParser)
60+
.postgresQueryParser(queryParser)
61+
.build();
62+
63+
String result = parser.parse(expression, context);
64+
65+
assertEquals("tags IS NULL", result);
66+
assertEquals(0, paramsBuilder.build().getObjectParams().size());
67+
}
68+
69+
@Test
70+
void testParseScalarRhsEq() {
71+
ArrayIdentifierExpression lhs = ArrayIdentifierExpression.of("tags", ArrayType.TEXT);
72+
ConstantExpression rhs = ConstantExpression.of("hygiene");
73+
RelationalExpression expression = RelationalExpression.of(lhs, EQ, rhs);
74+
75+
when(lhsParser.visit(any(ArrayIdentifierExpression.class))).thenReturn("tags");
76+
77+
String result = parser.parse(expression, context);
78+
79+
assertEquals("tags = ARRAY[?]::text[]", result);
80+
assertEquals(1, paramsBuilder.build().getObjectParams().size());
81+
assertEquals("hygiene", paramsBuilder.build().getObjectParams().get(1));
82+
}
83+
84+
@Test
85+
void testParseIterableRhsEq() {
86+
ArrayIdentifierExpression lhs = ArrayIdentifierExpression.of("tags", ArrayType.TEXT);
87+
ConstantExpression rhs = ConstantExpression.ofStrings(List.of("hygiene", "family-pack"));
88+
RelationalExpression expression = RelationalExpression.of(lhs, EQ, rhs);
89+
90+
when(lhsParser.visit(any(ArrayIdentifierExpression.class))).thenReturn("tags");
91+
92+
String result = parser.parse(expression, context);
93+
94+
assertEquals("tags = ARRAY[?, ?]::text[]", result);
95+
assertEquals(2, paramsBuilder.build().getObjectParams().size());
96+
assertEquals("hygiene", paramsBuilder.build().getObjectParams().get(1));
97+
assertEquals("family-pack", paramsBuilder.build().getObjectParams().get(2));
98+
}
99+
100+
@Test
101+
void testParseIterableRhsNeq() {
102+
ArrayIdentifierExpression lhs = ArrayIdentifierExpression.of("tags", ArrayType.TEXT);
103+
ConstantExpression rhs = ConstantExpression.ofStrings(List.of("a", "b"));
104+
RelationalExpression expression = RelationalExpression.of(lhs, NEQ, rhs);
105+
106+
when(lhsParser.visit(any(ArrayIdentifierExpression.class))).thenReturn("tags");
107+
108+
String result = parser.parse(expression, context);
109+
110+
assertEquals("tags != ARRAY[?, ?]::text[]", result);
111+
assertEquals(2, paramsBuilder.build().getObjectParams().size());
112+
assertEquals("a", paramsBuilder.build().getObjectParams().get(1));
113+
assertEquals("b", paramsBuilder.build().getObjectParams().get(2));
114+
}
115+
}

0 commit comments

Comments
 (0)