Skip to content

Commit 262e5c3

Browse files
pengpeng-lupengpeng-lu
authored
Support primitive type arrays contains (#3300)
This fixes #3674
1 parent 20f6da2 commit 262e5c3

File tree

10 files changed

+297
-89
lines changed

10 files changed

+297
-89
lines changed

‎fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/values/InOpValue.java‎

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -120,9 +120,6 @@ public Optional<QueryPredicate> toQueryPredicate(@Nullable final TypeRepository
120120
return compileTimeEvalMaybe(typeRepository);
121121
}
122122

123-
final var isLiteralList = inArrayValue.getCorrelatedTo().isEmpty();
124-
SemanticException.check(isLiteralList, SemanticException.ErrorCode.UNSUPPORTED);
125-
126123
if (typeRepository != null) {
127124
final var literalValue = Preconditions.checkNotNull(inArrayValue.evalWithoutStore(EvaluationContext.forTypeRepository(typeRepository)));
128125
return Optional.of(new ValuePredicate(probeValue, new Comparisons.ListComparison(Comparisons.Type.IN, (List<?>)literalValue)));
@@ -221,26 +218,31 @@ private static Value encapsulateInternal(@Nonnull final List<? extends Typed> ar
221218

222219
final Typed arg1 = arguments.get(1);
223220
final Type res1 = arg1.getResultType();
224-
SemanticException.check(res1.isArray(), SemanticException.ErrorCode.INCOMPATIBLE_TYPE);
221+
SemanticException.check(res1.isArray() || res0.isArray(), SemanticException.ErrorCode.INCOMPATIBLE_TYPE);
222+
223+
final Type arrayType = res0.isArray() ? res0 : res1;
224+
final Typed arrayArg = res0.isArray() ? arg0 : arg1;
225+
final Type operandType = res0.isArray() ? res1 : res0;
226+
final Typed operandArg = res0.isArray() ? arg1 : arg0;
225227

226-
final var arrayElementType = Objects.requireNonNull(((Type.Array) res1).getElementType());
227-
if (!arrayElementType.isUnresolved() && res0.getTypeCode() != arrayElementType.getTypeCode()) {
228-
final var maximumType = Type.maximumType(arg0.getResultType(), arrayElementType);
228+
final var arrayElementType = Objects.requireNonNull(((Type.Array) arrayType).getElementType());
229+
if (!arrayElementType.isUnresolved() && operandType.getTypeCode() != arrayElementType.getTypeCode()) {
230+
final var maximumType = Type.maximumType(operandType, arrayElementType);
229231
// Incompatible types
230232
SemanticException.check(maximumType != null, SemanticException.ErrorCode.INCOMPATIBLE_TYPE);
231233

232234
// Promote arg0 if the resulting type is different
233-
if (!arg0.getResultType().equals(maximumType)) {
234-
return new InOpValue(PromoteValue.inject((Value)arg0, maximumType), (Value)arg1);
235+
if (!operandType.equals(maximumType)) {
236+
return new InOpValue(PromoteValue.inject((Value)operandArg, maximumType), (Value)arrayArg);
235237
} else {
236-
return new InOpValue((Value)arg0, PromoteValue.inject((Value)arg1, new Type.Array(maximumType)));
238+
return new InOpValue((Value)operandArg, PromoteValue.inject((Value)arrayArg, new Type.Array(maximumType)));
237239
}
238240
}
239241

240-
if (res0.isRecord()) {
242+
if (operandType.isRecord()) {
241243
// we cannot yet promote this properly
242244
SemanticException.check(arrayElementType.isRecord(), SemanticException.ErrorCode.INCOMPATIBLE_TYPE);
243-
final var probeElementTypes = Objects.requireNonNull(((Type.Record)res0).getElementTypes());
245+
final var probeElementTypes = Objects.requireNonNull(((Type.Record)operandType).getElementTypes());
244246
final var inElementTypes = Objects.requireNonNull(((Type.Record)arrayElementType).getElementTypes());
245247
for (int i = 0; i < inElementTypes.size(); i++) {
246248
final var probeElementType = probeElementTypes.get(i);
@@ -251,7 +253,7 @@ private static Value encapsulateInternal(@Nonnull final List<? extends Typed> ar
251253
SemanticException.ErrorCode.INCOMPATIBLE_TYPE);
252254
}
253255
}
254-
return new InOpValue((Value)arg0, (Value)arg1);
256+
return new InOpValue((Value)operandArg, (Value)arrayArg);
255257
}
256258
}
257259

‎fdb-relational-core/src/main/antlr/RelationalParser.g4‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1162,6 +1162,7 @@ expressionAtom
11621162
inList
11631163
: '(' (queryExpressionBody | expressions) ')'
11641164
| preparedStatementParameter
1165+
| fullColumnName
11651166
;
11661167

11671168
preparedStatementParameter

‎fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizer.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -443,10 +443,16 @@ public Object visitPreparedStatementParameter(@Nonnull RelationalParser.Prepared
443443

444444
@Override
445445
public Object visitInPredicate(@Nonnull RelationalParser.InPredicateContext ctx) {
446+
if (ctx.NOT() != null) {
447+
ctx.NOT().accept(this);
448+
}
449+
446450
ctx.IN().accept(this);
447451

448452
if (ctx.inList().preparedStatementParameter() != null) {
449453
visit(ctx.inList().preparedStatementParameter());
454+
} else if (ctx.inList().fullColumnName() != null) {
455+
visit(ctx.inList().fullColumnName());
450456
} else {
451457
sqlCanonicalizer.append("( ");
452458
if (ParseHelpers.isConstant(ctx.inList().expressions())) {

‎fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/LogicalOperator.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,15 @@ private static LogicalOperator generateCorrelatedFieldAccess(@Nonnull Expression
263263
() -> String.format(Locale.ROOT, "join correlation can occur only on column of repeated type, not %s type", expression.getDataType()));
264264
final var explode = new ExplodeExpression(expression.getUnderlying());
265265
final var resultingQuantifier = Quantifier.forEach(Reference.initialOf(explode));
266-
final var outputAttributes = Expressions.of(convertToExpressions(resultingQuantifier));
266+
267+
Expressions outputAttributes;
268+
if (resultingQuantifier.getFlowedObjectType().isPrimitive()) {
269+
final ImmutableList.Builder<Expression> attributesBuilder = ImmutableList.builder();
270+
attributesBuilder.add(new Expression(alias, DataTypeUtils.toRelationalType(explode.getResultValue().getResultType()), resultingQuantifier.getFlowedObjectValue()));
271+
outputAttributes = Expressions.of(attributesBuilder.build());
272+
} else {
273+
outputAttributes = Expressions.of(convertToExpressions(resultingQuantifier));
274+
}
267275
return LogicalOperator.newOperator(alias, outputAttributes, resultingQuantifier);
268276
}
269277

‎fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/ExpressionVisitor.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -536,6 +536,10 @@ public Expression visitInList(@Nonnull RelationalParser.InListContext ctx) {
536536
final Expression result;
537537
if (ctx.preparedStatementParameter() != null) {
538538
result = visitPreparedStatementParameter(ctx.preparedStatementParameter());
539+
} else if (ctx.fullColumnName() != null) {
540+
result = visitFullColumnName(ctx.fullColumnName());
541+
// Validate that result is of type Array
542+
Assert.thatUnchecked(result.getDataType().getCode() == DataType.Code.ARRAY, ErrorCode.UNSUPPORTED_QUERY, "IN list with column reference must be of array type, but got: %s", result.getDataType().getCode());
539543
} else if (getDelegate().getPlanGenerationContext().shouldProcessLiteral() && ParseHelpers.isConstant(ctx.expressions())) {
540544
getDelegate().getPlanGenerationContext().startArrayLiteral();
541545
final var inListItems = visitExpressions(ctx.expressions());

‎fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizerTests.java‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1284,6 +1284,68 @@ void normalizeTemporarySqlFunctionStripsDefaultParameterLiterals() throws Except
12841284
Map.of(Options.Name.LOG_QUERY, false));
12851285
}
12861286

1287+
@Test
1288+
void visitInPredicateWithColumnReference() throws Exception {
1289+
// Test IN predicate with column reference instead of all constants
1290+
validate("select * from t1 where col1 in (fullColumnName)",
1291+
"select * from \"T1\" where \"COL1\" in ( \"FULLCOLUMNNAME\" ) ");
1292+
}
1293+
1294+
@Test
1295+
void visitInPredicateWithPreparedParameter() throws Exception {
1296+
// Test IN predicate with prepared parameter
1297+
java.sql.Array arrayParam = toArrayParameter(List.of("value1", "value2", "value3"));
1298+
validate("select * from t1 where col1 in ?",
1299+
PreparedParams.ofUnnamed(Map.of(1, arrayParam)),
1300+
"select * from \"T1\" where \"COL1\" in ? ",
1301+
Map.of(constantId(7), List.of("value1", "value2", "value3")));
1302+
}
1303+
1304+
@Test
1305+
void visitInPredicateWithMixedTypes() throws Exception {
1306+
// Test IN predicate with mixed constants and expressions
1307+
validate("select * from t1 where col1 in (10, col2 + 5, 'literal')",
1308+
"select * from \"T1\" where \"COL1\" in ( ? , \"COL2\" + ? , ? ) ",
1309+
Map.of(constantId(8), 10,
1310+
constantId(12), 5,
1311+
constantId(14), "literal"));
1312+
}
1313+
1314+
@Test
1315+
void visitInPredicateWithFullColumnNameInList() throws Exception {
1316+
// Test IN predicate with fullColumnName in the IN list - targets lines 453-454 in AstNormalizer
1317+
// This tests the case: 'apple' IN T.fruits where T.fruits is a column reference
1318+
validate("select * from T where 'apple' in T.fruits",
1319+
"select * from \"T\" where ? in \"T\" . \"FRUITS\" ",
1320+
Map.of(constantId(5), "apple"));
1321+
}
1322+
1323+
@Test
1324+
void visitInPredicateWithBooleanConstants() throws Exception {
1325+
// Test IN predicate with boolean constants
1326+
validate("select * from t1 where col1 in (true, false)",
1327+
"select * from \"T1\" where \"COL1\" in ( [ ] ) ",
1328+
Map.of(constantId(7), List.of(true, false)));
1329+
}
1330+
1331+
@Test
1332+
void visitNotInPredicateWithConstants() throws Exception {
1333+
// Test NOT IN predicate with constants - verifies that NOT token is handled correctly
1334+
// This tests the visitInPredicate method's handling of NOT token in AstNormalizer
1335+
validate("select * from t1 where col1 not in (10, 20, 30)",
1336+
"select * from \"T1\" where \"COL1\" not in ( [ ] ) ",
1337+
Map.of(constantId(8), List.of(10, 20, 30)));
1338+
}
1339+
1340+
@Test
1341+
void visitNotInPredicateWithColumnReference() throws Exception {
1342+
// Test NOT IN predicate with column reference - this tests lines 453-454 in AstNormalizer
1343+
// where ctx.inList().fullColumnName() != null for a NOT IN predicate
1344+
validate("select * from T where 'apple' not in T.fruits",
1345+
"select * from \"T\" where ? not in \"T\" . \"FRUITS\" ",
1346+
Map.of(constantId(5), "apple"));
1347+
}
1348+
12871349
@Nonnull
12881350
private String normalizeQuery(@Nonnull final String functionDdl) throws RelationalException {
12891351
final var normalizer = AstNormalizer.normalizeAst(fakeSchemaTemplate,

‎fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/StandardQueryTests.java‎

Lines changed: 63 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,6 @@ void simpleSelect() throws Exception {
159159

160160
@Test
161161
void simpleSelectWithNonNullableArrays() throws Exception {
162-
// var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplateWithNonNullableArrays).build();
163162
try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplateWithNonNullableArrays).build()) {
164163
try (var statement = ddl.setSchemaAndGetConnection().createStatement()) {
165164
var insertedRecord = insertRestaurantComplexRecord(statement);
@@ -1362,6 +1361,69 @@ void unionIsNotSupported() throws Exception {
13621361
}
13631362
}
13641363

1364+
@Test
1365+
void structArrayContains() throws Exception {
1366+
final String schemaTemplate = "CREATE TYPE AS STRUCT A(col2 string, col3 bigint, col4 bigint) " +
1367+
"CREATE TABLE T1(col1 bigint, a A Array, col5 bigint, primary key(col1))";
1368+
try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).build()) {
1369+
try (var statement = ddl.setSchemaAndGetConnection().createStatement()) {
1370+
statement.executeUpdate("insert into t1 values (42, [('Apple', 1, 100), ('Orange', 2, 200)], 142), (44, [('Grape', 3, 300), ('Pear', 4, 400)], 144)");
1371+
Assertions.assertTrue(statement.execute("SELECT T1.col5 FROM T1 where exists (SELECT 1 FROM T1.A r where r.col2 = 'Grape') "));
1372+
try (final RelationalResultSet resultSet = statement.getResultSet()) {
1373+
ResultSetAssert.assertThat(resultSet).hasNextRow()
1374+
.isRowExactly(144L)
1375+
.hasNoNextRow();
1376+
}
1377+
// another way of query
1378+
Assertions.assertTrue(statement.execute("SELECT T1.col5 FROM T1, (SELECT col2, col3 FROM T1.A) X where X.col2 = 'Grape'"));
1379+
try (final RelationalResultSet resultSet = statement.getResultSet()) {
1380+
ResultSetAssert.assertThat(resultSet).hasNextRow()
1381+
.isRowExactly(144L)
1382+
.hasNoNextRow();
1383+
}
1384+
Assertions.assertTrue(statement.execute("SELECT T1.col5 FROM T1 where exists (SELECT 1 FROM T1.A r where r.col2 in ('Grape', 'Orange'))"));
1385+
try (final RelationalResultSet resultSet = statement.getResultSet()) {
1386+
ResultSetAssert.assertThat(resultSet).hasNextRow()
1387+
.isRowExactly(142L)
1388+
.hasNextRow()
1389+
.isRowExactly(144L)
1390+
.hasNoNextRow();
1391+
}
1392+
}
1393+
}
1394+
}
1395+
1396+
@Test
1397+
void primitiveArrayContains() throws Exception {
1398+
final String schemaTemplate = "CREATE TABLE T1(col1 bigint, a string Array, primary key(col1))";
1399+
try (var ddl = Ddl.builder().database(URI.create("/TEST/QT")).relationalExtension(relationalExtension).schemaTemplate(schemaTemplate).build()) {
1400+
try (var statement = ddl.setSchemaAndGetConnection().createStatement()) {
1401+
statement.executeUpdate("insert into t1 values (42, ['Apple', 'Orange']), (44, ['Grape', 'Pear'])");
1402+
Assertions.assertTrue(statement.execute("SELECT * FROM T1 where exists (SELECT 1 FROM T1.A r where r = 'Grape')"));
1403+
try (final RelationalResultSet resultSet = statement.getResultSet()) {
1404+
ResultSetAssert.assertThat(resultSet).hasNextRow()
1405+
.isRowExactly(44L, EmbeddedRelationalArray.newBuilder().addString("Grape").addString("Pear").build())
1406+
.hasNoNextRow();
1407+
}
1408+
Assertions.assertTrue(statement.execute("SELECT * FROM T1 where exists (SELECT 1 FROM T1.A r where r in ('Grape', 'Orange'))"));
1409+
try (final RelationalResultSet resultSet = statement.getResultSet()) {
1410+
ResultSetAssert.assertThat(resultSet).hasNextRow()
1411+
.isRowExactly(42L, EmbeddedRelationalArray.newBuilder().addString("Apple").addString("Orange").build())
1412+
.hasNextRow()
1413+
.isRowExactly(44L, EmbeddedRelationalArray.newBuilder().addString("Grape").addString("Pear").build())
1414+
.hasNoNextRow();
1415+
}
1416+
1417+
Assertions.assertTrue(statement.execute("SELECT * FROM T1 where 'Grape' in a"));
1418+
try (final RelationalResultSet resultSet = statement.getResultSet()) {
1419+
ResultSetAssert.assertThat(resultSet).hasNextRow()
1420+
.isRowExactly(44L, EmbeddedRelationalArray.newBuilder().addString("Grape").addString("Pear").build())
1421+
.hasNoNextRow();
1422+
}
1423+
}
1424+
}
1425+
}
1426+
13651427
@Test
13661428
void cteWorksCorrectly() throws Exception {
13671429
final String schemaTemplate = "CREATE TABLE T1(pk bigint, a bigint, b bigint, c bigint, PRIMARY KEY(pk))";

0 commit comments

Comments
 (0)