Skip to content

Commit b7f7b87

Browse files
yuancuyuancu
committed
Reorganize class locations
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 53a4118 commit b7f7b87

13 files changed

Lines changed: 1630 additions & 243 deletions

File tree

‎core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java‎

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import java.util.function.BiFunction;
1818
import lombok.Getter;
1919
import lombok.Setter;
20+
import org.apache.calcite.config.NullCollation;
2021
import org.apache.calcite.rex.RexCorrelVariable;
2122
import org.apache.calcite.rex.RexLambdaRef;
2223
import org.apache.calcite.rex.RexNode;
@@ -26,8 +27,11 @@
2627
import org.apache.calcite.tools.RelBuilder;
2728
import org.opensearch.sql.ast.expression.UnresolvedExpression;
2829
import org.opensearch.sql.calcite.utils.CalciteToolsHelper;
30+
import org.opensearch.sql.calcite.validate.OpenSearchSparkSqlDialect;
31+
import org.opensearch.sql.calcite.validate.PplTypeCoercion;
32+
import org.opensearch.sql.calcite.validate.PplTypeCoercionRule;
33+
import org.opensearch.sql.calcite.validate.PplValidator;
2934
import org.opensearch.sql.calcite.validate.SqlOperatorTableProvider;
30-
import org.opensearch.sql.calcite.validate.TypeChecker;
3135
import org.opensearch.sql.common.setting.Settings;
3236
import org.opensearch.sql.executor.QueryType;
3337
import org.opensearch.sql.expression.function.FunctionProperties;
@@ -105,8 +109,21 @@ public SqlValidator getValidator() {
105109
throw new IllegalStateException(
106110
"SqlOperatorTableProvider must be set before creating CalcitePlanContext");
107111
}
112+
SqlValidator.Config validatorConfig =
113+
SqlValidator.Config.DEFAULT
114+
.withTypeCoercionRules(PplTypeCoercionRule.instance())
115+
.withTypeCoercionFactory(PplTypeCoercion::create)
116+
// Use lenient conformance for PPL compatibility
117+
.withConformance(OpenSearchSparkSqlDialect.DEFAULT.getConformance())
118+
// Use Spark SQL's NULL collation (NULLs sorted LOW/FIRST)
119+
.withDefaultNullCollation(NullCollation.LOW);
108120
validator =
109-
TypeChecker.getValidator(statement, config, operatorTableProvider.getOperatorTable());
121+
PplValidator.create(
122+
statement,
123+
config,
124+
operatorTableProvider.getOperatorTable(),
125+
TYPE_FACTORY,
126+
validatorConfig);
110127
}
111128
return validator;
112129
}

‎core/src/main/java/org/opensearch/sql/calcite/validate/PplTypeCoercion.java‎

Lines changed: 10 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,6 @@
88
import static java.util.Objects.requireNonNull;
99
import static org.opensearch.sql.calcite.validate.ValidationUtils.createUDTWithAttributes;
1010

11-
import java.util.List;
12-
import java.util.Map;
13-
import java.util.Set;
14-
import java.util.stream.IntStream;
1511
import org.apache.calcite.adapter.java.JavaTypeFactory;
1612
import org.apache.calcite.rel.type.RelDataType;
1713
import org.apache.calcite.rel.type.RelDataTypeFactory;
@@ -27,6 +23,7 @@
2723
import org.apache.calcite.sql.type.SqlTypeUtil;
2824
import org.apache.calcite.sql.validate.SqlValidator;
2925
import org.apache.calcite.sql.validate.SqlValidatorScope;
26+
import org.apache.calcite.sql.validate.implicit.TypeCoercion;
3027
import org.apache.calcite.sql.validate.implicit.TypeCoercionImpl;
3128
import org.checkerframework.checker.nullness.qual.Nullable;
3229
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
@@ -42,49 +39,22 @@
4239
* are not allowed in PPL semantics.
4340
*/
4441
public class PplTypeCoercion extends TypeCoercionImpl {
45-
// A blacklist of coercions that are not allowed in PPL.
46-
// key cannot be cast from values
47-
private static final Map<SqlTypeFamily, Set<SqlTypeFamily>> BLACKLISTED_COERCIONS;
4842

49-
static {
50-
// Initialize the blacklist for coercions that are not allowed in PPL.
51-
BLACKLISTED_COERCIONS = Map.of();
43+
/**
44+
* Creates a custom TypeCoercion instance for PPL. This can be used as a TypeCoercionFactory.
45+
*
46+
* @param typeFactory the type factory
47+
* @param validator the SQL validator
48+
* @return custom PplTypeCoercion instance
49+
*/
50+
public static TypeCoercion create(RelDataTypeFactory typeFactory, SqlValidator validator) {
51+
return new PplTypeCoercion(typeFactory, validator);
5252
}
5353

5454
public PplTypeCoercion(RelDataTypeFactory typeFactory, SqlValidator validator) {
5555
super(typeFactory, validator);
5656
}
5757

58-
@Override
59-
public boolean builtinFunctionCoercion(
60-
SqlCallBinding binding,
61-
List<RelDataType> operandTypes,
62-
List<SqlTypeFamily> expectedFamilies) {
63-
assert binding.getOperandCount() == operandTypes.size();
64-
if (IntStream.range(0, operandTypes.size())
65-
.anyMatch(i -> isBlacklistedCoercion(operandTypes.get(i), expectedFamilies.get(i)))) {
66-
return false;
67-
}
68-
return super.builtinFunctionCoercion(binding, operandTypes, expectedFamilies);
69-
}
70-
71-
/**
72-
* Checks if a type coercion is blacklisted based on PPL rules.
73-
*
74-
* @param operandType the actual type of the operand
75-
* @param expectedFamily the expected type family
76-
* @return true if the coercion is blacklisted, false otherwise
77-
*/
78-
private boolean isBlacklistedCoercion(RelDataType operandType, SqlTypeFamily expectedFamily) {
79-
if (BLACKLISTED_COERCIONS.containsKey(expectedFamily)) {
80-
Set<SqlTypeFamily> blacklistedFamilies = BLACKLISTED_COERCIONS.get(expectedFamily);
81-
if (blacklistedFamilies.contains(operandType.getSqlTypeName().getFamily())) {
82-
return true;
83-
}
84-
}
85-
return false;
86-
}
87-
8858
@Override
8959
public @Nullable RelDataType implicitCast(RelDataType in, SqlTypeFamily expected) {
9060
RelDataType casted = super.implicitCast(in, expected);

‎core/src/main/java/org/opensearch/sql/calcite/validate/PplValidator.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,19 +9,26 @@
99

1010
import java.util.List;
1111
import java.util.function.Function;
12+
import org.apache.calcite.jdbc.CalcitePrepare;
13+
import org.apache.calcite.jdbc.CalciteSchema;
14+
import org.apache.calcite.prepare.CalciteCatalogReader;
1215
import org.apache.calcite.rel.type.RelDataType;
1316
import org.apache.calcite.rel.type.RelDataTypeFactory;
1417
import org.apache.calcite.rel.type.RelDataTypeField;
1518
import org.apache.calcite.rel.type.RelRecordType;
19+
import org.apache.calcite.schema.SchemaPlus;
20+
import org.apache.calcite.server.CalciteServerStatement;
1621
import org.apache.calcite.sql.SqlNode;
1722
import org.apache.calcite.sql.SqlOperatorTable;
1823
import org.apache.calcite.sql.type.ArraySqlType;
1924
import org.apache.calcite.sql.type.MapSqlType;
2025
import org.apache.calcite.sql.type.MultisetSqlType;
2126
import org.apache.calcite.sql.type.SqlTypeName;
27+
import org.apache.calcite.sql.validate.SqlValidator;
2228
import org.apache.calcite.sql.validate.SqlValidatorCatalogReader;
2329
import org.apache.calcite.sql.validate.SqlValidatorImpl;
2430
import org.apache.calcite.sql.validate.SqlValidatorScope;
31+
import org.apache.calcite.tools.FrameworkConfig;
2532
import org.checkerframework.checker.nullness.qual.Nullable;
2633
import org.opensearch.sql.calcite.type.AbstractExprRelDataType;
2734
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
@@ -45,6 +52,31 @@ public class PplValidator extends SqlValidatorImpl {
4552
*/
4653
private boolean top;
4754

55+
/**
56+
* Creates a SqlValidator configured for PPL validation.
57+
*
58+
* @param statement Calcite server statement
59+
* @param frameworkConfig Framework configuration
60+
* @param operatorTable SQL operator table to use for validation
61+
* @return configured SqlValidator instance
62+
*/
63+
public static PplValidator create(
64+
CalciteServerStatement statement,
65+
FrameworkConfig frameworkConfig,
66+
SqlOperatorTable operatorTable,
67+
RelDataTypeFactory typeFactory,
68+
SqlValidator.Config validatorConfig) {
69+
SchemaPlus defaultSchema = frameworkConfig.getDefaultSchema();
70+
71+
final CalcitePrepare.Context prepareContext = statement.createPrepareContext();
72+
final CalciteSchema schema =
73+
defaultSchema != null ? CalciteSchema.from(defaultSchema) : prepareContext.getRootSchema();
74+
CalciteCatalogReader catalogReader =
75+
new CalciteCatalogReader(
76+
schema.root(), schema.path(null), typeFactory, prepareContext.config());
77+
return new PplValidator(operatorTable, catalogReader, typeFactory, validatorConfig);
78+
}
79+
4880
/**
4981
* Creates a PPL validator.
5082
*

‎core/src/main/java/org/opensearch/sql/calcite/validate/TypeChecker.java‎

Lines changed: 0 additions & 83 deletions
This file was deleted.

core/src/main/java/org/opensearch/sql/calcite/validate/PplRelToSqlNodeConverter.java renamed to core/src/main/java/org/opensearch/sql/calcite/validate/converters/PplRelToSqlNodeConverter.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
package org.opensearch.sql.calcite.validate;
6+
package org.opensearch.sql.calcite.validate.converters;
77

88
import org.apache.calcite.rel.core.Correlate;
99
import org.apache.calcite.rel.core.Join;

core/src/main/java/org/opensearch/sql/calcite/validate/PplSqlToRelConverter.java renamed to core/src/main/java/org/opensearch/sql/calcite/validate/converters/PplSqlToRelConverter.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
package org.opensearch.sql.calcite.validate;
6+
package org.opensearch.sql.calcite.validate.converters;
77

88
import java.util.List;
99
import org.apache.calcite.plan.RelOptCluster;

‎core/src/main/java/org/opensearch/sql/calcite/validate/shuttles/SkipRelValidationShuttle.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,9 @@
7878
* <p><b>Note for developers</b>: when validations fail during developing new features, please try
7979
* to solve the root cause instead of adding skipping rules here. Under rare cases when you have to
8080
* skip validation, please document the exact reason.
81+
*
82+
* <p><b>WARNING</b>: When a skip pattern is detected, we bypass the entire validation pipeline,
83+
* skipping potentially useful transformation relying on rewriting SQL node
8184
*/
8285
public class SkipRelValidationShuttle extends RelShuttleImpl {
8386
private boolean shouldSkip = false;
@@ -93,6 +96,8 @@ public class SkipRelValidationShuttle extends RelShuttleImpl {
9396
public static final List<Predicate<LogicalValues>> SKIP_VALUES;
9497

9598
static {
99+
// TODO: Make incompatible operations like bin-on-timestamp a validatable UDFs so that they can
100+
// be still be converted to SqlNode and back to RelNode
96101
Predicate<RexCall> binOnTimestamp =
97102
call -> {
98103
if ("WIDTH_BUCKET".equalsIgnoreCase(call.getOperator().getName())) {

‎core/src/main/java/org/opensearch/sql/calcite/validate/shuttles/SqlRewriteShuttle.java‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.validate.shuttles;
7+
8+
import java.util.Collections;
9+
import java.util.List;
10+
import org.apache.calcite.sql.SqlBasicCall;
11+
import org.apache.calcite.sql.SqlCall;
12+
import org.apache.calcite.sql.SqlIdentifier;
13+
import org.apache.calcite.sql.SqlKind;
14+
import org.apache.calcite.sql.SqlNode;
15+
import org.apache.calcite.sql.SqlNodeList;
16+
import org.apache.calcite.sql.fun.SqlCountAggFunction;
17+
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
18+
import org.apache.calcite.sql.util.SqlShuttle;
19+
import org.opensearch.sql.calcite.OpenSearchSchema;
20+
21+
public class SqlRewriteShuttle extends SqlShuttle {
22+
@Override
23+
public SqlNode visit(SqlIdentifier id) {
24+
// Remove database qualifier, keeping only table name
25+
if (id.names.size() == 2 && OpenSearchSchema.OPEN_SEARCH_SCHEMA_NAME.equals(id.names.get(0))) {
26+
return new SqlIdentifier(Collections.singletonList(id.names.get(1)), id.getParserPosition());
27+
}
28+
return id;
29+
}
30+
31+
@Override
32+
public @org.checkerframework.checker.nullness.qual.Nullable SqlNode visit(SqlCall call) {
33+
if (call.getOperator() instanceof SqlCountAggFunction && call.getOperandList().isEmpty()) {
34+
// Convert COUNT() to COUNT(*) so that SqlCall.isCountStar() resolves to True
35+
// This is useful when deriving the return types in SqlCountAggFunction#deriveType
36+
call =
37+
new SqlBasicCall(
38+
SqlStdOperatorTable.COUNT,
39+
List.of(SqlIdentifier.STAR),
40+
call.getParserPosition(),
41+
call.getFunctionQuantifier());
42+
} else if (call.getKind() == SqlKind.IN || call.getKind() == SqlKind.NOT_IN) {
43+
// Fix for tuple IN / NOT IN queries: Convert SqlNodeList to ROW SqlCall
44+
//
45+
// When RelToSqlConverter converts a tuple expression like (id, name) back to
46+
// SqlNode, it generates a bare SqlNodeList instead of wrapping it in a ROW
47+
// operator. This causes validation to fail because:
48+
// 1. SqlValidator.deriveType() doesn't know how to handle SqlNodeList
49+
// 2. SqlToRelConverter.visit(SqlNodeList) throws UnsupportedOperationException
50+
//
51+
// For example, the query:
52+
// WHERE (id, name) NOT IN (SELECT uid, name FROM ...)
53+
//
54+
// After Rel-to-SQL conversion becomes:
55+
// IN operator with operands: [SqlNodeList[id, name], SqlSelect[...]]
56+
//
57+
// But it should be:
58+
// IN operator with operands: [ROW(id, name), SqlSelect[...]]
59+
//
60+
// This fix wraps the SqlNodeList in a ROW SqlCall before validation,
61+
// ensuring proper type derivation and subsequent SQL-to-Rel conversion.
62+
if (!call.getOperandList().isEmpty()
63+
&& call.getOperandList().get(0) instanceof SqlNodeList nodes) {
64+
call.setOperand(0, SqlStdOperatorTable.ROW.createCall(nodes));
65+
}
66+
}
67+
return super.visit(call);
68+
}
69+
}

0 commit comments

Comments
 (0)