Skip to content

Commit 87227be

Browse files
yuancuyuancu
committed
Enable IP comparison (1332/1559 | 1409/1915)
- allow type cast - rewrite call to sql compare to custom ip comapre Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 2e6c61c commit 87227be

12 files changed

Lines changed: 298 additions & 152 deletions

File tree

‎core/src/main/java/org/opensearch/sql/calcite/utils/OpenSearchTypeFactory.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,12 +42,14 @@
4242
import org.apache.calcite.sql.SqlCollation;
4343
import org.apache.calcite.sql.type.SqlTypeName;
4444
import org.apache.calcite.sql.type.SqlTypeUtil;
45+
import org.checkerframework.checker.nullness.qual.Nullable;
4546
import org.opensearch.sql.calcite.type.AbstractExprRelDataType;
4647
import org.opensearch.sql.calcite.type.ExprBinaryType;
4748
import org.opensearch.sql.calcite.type.ExprDateType;
4849
import org.opensearch.sql.calcite.type.ExprIPType;
4950
import org.opensearch.sql.calcite.type.ExprTimeStampType;
5051
import org.opensearch.sql.calcite.type.ExprTimeType;
52+
import org.opensearch.sql.calcite.validate.PplTypeCoercionRule;
5153
import org.opensearch.sql.data.model.ExprValue;
5254
import org.opensearch.sql.data.model.ExprValueUtils;
5355
import org.opensearch.sql.data.type.ExprCoreType;
@@ -327,6 +329,12 @@ public Type getJavaClass(RelDataType type) {
327329
return super.getJavaClass(type);
328330
}
329331

332+
@Override
333+
public @Nullable RelDataType leastRestrictive(List<RelDataType> types) {
334+
// In parent: leastRestrictive(types, SqlTypeMappingRules.instance(false))
335+
return leastRestrictive(types, PplTypeCoercionRule.assignmentInstance());
336+
}
337+
330338
/**
331339
* Whether a given RelDataType is a user-defined type (UDT)
332340
*
@@ -436,4 +444,11 @@ public static boolean isTime(RelDataType type) {
436444
public static boolean isCharacter(RelDataType type) {
437445
return !isUserDefinedType(type) && SqlTypeUtil.isCharacter(type);
438446
}
447+
448+
public static boolean isIp(RelDataType type) {
449+
if (isUserDefinedType(type)) {
450+
return ((AbstractExprRelDataType<?>) type).getUdt() == ExprUDT.EXPR_IP;
451+
}
452+
return false;
453+
}
439454
}

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

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.validate;
7+
8+
import java.util.HashMap;
9+
import java.util.Map;
10+
import org.apache.calcite.rel.type.RelDataType;
11+
import org.apache.calcite.rex.RexCall;
12+
import org.apache.calcite.sql.SqlCall;
13+
import org.apache.calcite.sql.SqlFunction;
14+
import org.apache.calcite.sql.SqlOperator;
15+
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
16+
import org.apache.calcite.sql2rel.ReflectiveConvertletTable;
17+
import org.apache.calcite.sql2rel.SqlRexConvertlet;
18+
import org.apache.calcite.sql2rel.StandardConvertletTable;
19+
import org.checkerframework.checker.initialization.qual.UnderInitialization;
20+
import org.checkerframework.checker.nullness.qual.Nullable;
21+
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
22+
import org.opensearch.sql.expression.function.PPLBuiltinOperators;
23+
24+
public class PplConvertletTable extends ReflectiveConvertletTable {
25+
public static PplConvertletTable INSTANCE = new PplConvertletTable();
26+
private final Map<SqlOperator, SqlRexConvertlet> map = new HashMap<>();
27+
28+
private PplConvertletTable() {
29+
super();
30+
registerOperator(SqlStdOperatorTable.EQUALS, ipConvertlet(PPLBuiltinOperators.EQUALS_IP));
31+
registerOperator(
32+
SqlStdOperatorTable.NOT_EQUALS, ipConvertlet(PPLBuiltinOperators.NOT_EQUALS_IP));
33+
registerOperator(
34+
SqlStdOperatorTable.GREATER_THAN, ipConvertlet(PPLBuiltinOperators.GREATER_IP));
35+
registerOperator(
36+
SqlStdOperatorTable.GREATER_THAN_OR_EQUAL, ipConvertlet(PPLBuiltinOperators.GTE_IP));
37+
registerOperator(SqlStdOperatorTable.LESS_THAN, ipConvertlet(PPLBuiltinOperators.LESS_IP));
38+
registerOperator(
39+
SqlStdOperatorTable.LESS_THAN_OR_EQUAL, ipConvertlet(PPLBuiltinOperators.LTE_IP));
40+
}
41+
42+
@Override
43+
public @Nullable SqlRexConvertlet get(SqlCall call) {
44+
SqlRexConvertlet custom = map.get(call.getOperator());
45+
if (custom != null) return custom;
46+
return StandardConvertletTable.INSTANCE.get(call);
47+
}
48+
49+
/** Registers a convertlet for a given operator instance. */
50+
private void registerOperator(
51+
@UnderInitialization PplConvertletTable this, SqlOperator op, SqlRexConvertlet convertlet) {
52+
map.put(op, convertlet);
53+
}
54+
55+
private SqlRexConvertlet ipConvertlet(SqlFunction substitute) {
56+
return (cx, call) -> {
57+
final RexCall e = (RexCall) StandardConvertletTable.INSTANCE.convertCall(cx, call);
58+
RelDataType type1 = e.getOperands().get(0).getType();
59+
RelDataType type2 = e.getOperands().get(1).getType();
60+
if (OpenSearchTypeFactory.isIp(type1) || OpenSearchTypeFactory.isIp(type2)) {
61+
return StandardConvertletTable.INSTANCE.convertFunction(cx, substitute, call);
62+
}
63+
return e;
64+
};
65+
}
66+
}

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,10 @@
1010
import org.apache.calcite.rel.RelNode;
1111
import org.apache.calcite.rel.RelShuttleImpl;
1212
import org.apache.calcite.rex.RexBuilder;
13-
import org.apache.calcite.rex.RexCall;
1413
import org.apache.calcite.rex.RexLiteral;
1514
import org.apache.calcite.rex.RexNode;
1615
import org.apache.calcite.rex.RexShuttle;
1716
import org.apache.calcite.sql.SqlIntervalQualifier;
18-
import org.apache.calcite.sql.SqlKind;
1917

2018
/**
2119
* A RelShuttle that recursively visits all RelNodes and their RexNode expressions to fix interval

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import org.apache.calcite.sql.SqlDynamicParam;
2222
import org.apache.calcite.sql.SqlNode;
2323
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
24-
import org.apache.calcite.sql.type.SqlTypeCoercionRule;
2524
import org.apache.calcite.sql.type.SqlTypeFamily;
2625
import org.apache.calcite.sql.type.SqlTypeMappingRule;
2726
import org.apache.calcite.sql.type.SqlTypeName;
@@ -50,13 +49,6 @@ public class PplTypeCoercion extends TypeCoercionImpl {
5049
static {
5150
// Initialize the blacklist for coercions that are not allowed in PPL.
5251
BLACKLISTED_COERCIONS = Map.of();
53-
// Map.of(
54-
// SqlTypeFamily.CHARACTER,
55-
// Set.of(SqlTypeFamily.NUMERIC),
56-
// SqlTypeFamily.STRING,
57-
// Set.of(SqlTypeFamily.NUMERIC),
58-
// SqlTypeFamily.NUMERIC,
59-
// Set.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.STRING));
6052
}
6153

6254
public PplTypeCoercion(RelDataTypeFactory typeFactory, SqlValidator validator) {
@@ -153,6 +145,7 @@ protected boolean dateTimeStringEquality(
153145
// Prepend following rules for datetime comparisons:
154146
// - (date, time) -> timestamp
155147
// - (time, timestamp) -> timestamp
148+
// - (ip, string) -> ip
156149
if (type1 != null & type2 != null) {
157150
boolean anyNullable = type1.isNullable() || type2.isNullable();
158151
if ((SqlTypeUtil.isDate(type1) && OpenSearchTypeFactory.isTime(type2))
@@ -166,6 +159,12 @@ protected boolean dateTimeStringEquality(
166159
if (SqlTypeUtil.isTimestamp(type1) && OpenSearchTypeFactory.isTime(type2)) {
167160
return factory.createTypeWithNullability(type1, anyNullable);
168161
}
162+
if (OpenSearchTypeFactory.isIp(type1) && OpenSearchTypeFactory.isCharacter(type2)) {
163+
return factory.createTypeWithNullability(type1, anyNullable);
164+
}
165+
if (OpenSearchTypeFactory.isCharacter(type1) && OpenSearchTypeFactory.isIp(type2)) {
166+
return factory.createTypeWithNullability(type2, anyNullable);
167+
}
169168
}
170169
return super.commonTypeForBinaryComparison(type1, type2);
171170
}
@@ -194,7 +193,7 @@ protected boolean coerceOperandType(
194193
}
195194

196195
// Check it early.
197-
if (!needToCast(scope, operand, targetType, SqlTypeCoercionRule.lenientInstance())) {
196+
if (!needToCast(scope, operand, targetType, PplTypeCoercionRule.lenientInstance())) {
198197
return false;
199198
}
200199
// Fix up nullable attr.
@@ -206,7 +205,7 @@ protected boolean coerceOperandType(
206205
}
207206

208207
private static SqlNode castTo(SqlNode node, RelDataType type) {
209-
if (OpenSearchTypeFactory.isDatetime(type)) {
208+
if (OpenSearchTypeFactory.isDatetime(type) || OpenSearchTypeFactory.isIp(type)) {
210209
ExprType exprType = OpenSearchTypeFactory.convertRelDataTypeToExprType(type);
211210
return switch (exprType) {
212211
case ExprCoreType.DATE ->
@@ -215,6 +214,7 @@ private static SqlNode castTo(SqlNode node, RelDataType type) {
215214
PPLBuiltinOperators.TIMESTAMP.createCall(node.getParserPosition(), node);
216215
case ExprCoreType.TIME ->
217216
PPLBuiltinOperators.TIME.createCall(node.getParserPosition(), node);
217+
case ExprCoreType.IP -> PPLBuiltinOperators.IP.createCall(node.getParserPosition(), node);
218218
default -> throw new UnsupportedOperationException("Unsupported type: " + exprType);
219219
};
220220
}

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

Lines changed: 63 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,69 @@
55

66
package org.opensearch.sql.calcite.validate;
77

8+
import com.google.common.collect.ImmutableMap;
9+
import com.google.common.collect.ImmutableSet;
10+
import java.util.HashSet;
11+
import java.util.Map;
12+
import java.util.Set;
13+
import java.util.stream.Collectors;
14+
import java.util.stream.Stream;
15+
import lombok.NonNull;
16+
import org.apache.calcite.sql.type.SqlTypeAssignmentRule;
817
import org.apache.calcite.sql.type.SqlTypeCoercionRule;
18+
import org.apache.calcite.sql.type.SqlTypeName;
919

10-
public class PplTypeCoercionRules extends SqlTypeCoercionRule {
20+
public class PplTypeCoercionRule {
21+
/**
22+
* PPL-specific additional type mapping rules
23+
*
24+
* <ul>
25+
* <li>IP -> IP
26+
* <li>CHARACTER -> IP
27+
* </ul>
28+
*/
29+
private static final Map<SqlTypeName, ImmutableSet<@NonNull SqlTypeName>> additionalMapping =
30+
Map.of(
31+
SqlTypeName.OTHER,
32+
ImmutableSet.of(SqlTypeName.OTHER, SqlTypeName.VARCHAR, SqlTypeName.CHAR));
33+
34+
private static final SqlTypeCoercionRule INSTANCE =
35+
SqlTypeCoercionRule.instance(
36+
mergeMapping(SqlTypeCoercionRule.instance().getTypeMapping(), additionalMapping));
37+
private static final SqlTypeCoercionRule LENIENT_INSTANCE =
38+
SqlTypeCoercionRule.instance(
39+
mergeMapping(SqlTypeCoercionRule.lenientInstance().getTypeMapping(), additionalMapping));
40+
private static final SqlTypeCoercionRule ASSIGNMENT_INSTANCE =
41+
SqlTypeCoercionRule.instance(
42+
mergeMapping(SqlTypeAssignmentRule.instance().getTypeMapping(), additionalMapping));
43+
44+
public static SqlTypeCoercionRule instance() {
45+
return INSTANCE;
46+
}
47+
48+
/** Returns an instance that allows more lenient type coercion. */
49+
public static SqlTypeCoercionRule lenientInstance() {
50+
return LENIENT_INSTANCE;
51+
}
52+
53+
/** Rules that determine whether a type is assignable from another type. */
54+
public static SqlTypeCoercionRule assignmentInstance() {
55+
return ASSIGNMENT_INSTANCE;
56+
}
57+
58+
private static <T> Map<T, ImmutableSet<@NonNull T>> mergeMapping(
59+
Map<T, ImmutableSet<@NonNull T>> base, Map<T, ImmutableSet<@NonNull T>> addition) {
60+
return Stream.concat(base.entrySet().stream(), addition.entrySet().stream())
61+
.collect(
62+
Collectors.collectingAndThen(
63+
Collectors.toMap(
64+
Map.Entry::getKey,
65+
Map.Entry::getValue,
66+
(b, a) -> {
67+
Set<T> combined = new HashSet<>(b);
68+
combined.addAll(a);
69+
return ImmutableSet.copyOf(combined);
70+
}),
71+
ImmutableMap::copyOf));
72+
}
1173
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
* PPL-specific operator tables.
2727
*/
2828
public class TypeChecker {
29+
2930
/**
3031
* Creates a SqlValidator configured for PPL validation.
3132
*
@@ -63,8 +64,7 @@ public static SqlValidator getValidator(
6364
* @return SqlTypeCoercionRule instance
6465
*/
6566
public static SqlTypeCoercionRule getTypeCoercionRule() {
66-
var defaultMapping = SqlTypeCoercionRule.instance().getTypeMapping();
67-
return SqlTypeCoercionRule.instance(defaultMapping);
67+
return PplTypeCoercionRule.instance();
6868
}
6969

7070
/**

‎core/src/main/java/org/opensearch/sql/executor/QueryService.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,6 @@
4040
import org.apache.calcite.sql.util.SqlShuttle;
4141
import org.apache.calcite.sql.validate.SqlValidator;
4242
import org.apache.calcite.sql2rel.SqlToRelConverter;
43-
import org.apache.calcite.sql2rel.StandardConvertletTable;
4443
import org.apache.calcite.tools.FrameworkConfig;
4544
import org.apache.calcite.tools.Frameworks;
4645
import org.apache.calcite.tools.Programs;
@@ -54,6 +53,7 @@
5453
import org.opensearch.sql.calcite.SysLimit;
5554
import org.opensearch.sql.calcite.plan.LogicalSystemLimit;
5655
import org.opensearch.sql.calcite.plan.LogicalSystemLimit.SystemLimitType;
56+
import org.opensearch.sql.calcite.validate.PplConvertletTable;
5757
import org.opensearch.sql.calcite.validate.PplRelToSqlNodeConverter;
5858
import org.opensearch.sql.calcite.validate.PplRelToSqlRelShuttle;
5959
import org.opensearch.sql.common.response.ResponseListener;
@@ -81,7 +81,7 @@ public class QueryService {
8181
private final Planner planner;
8282
private DataSourceService dataSourceService;
8383
private Settings settings;
84-
private static final PplRelToSqlNodeConverter converter =
84+
private static final PplRelToSqlNodeConverter rel2sql =
8585
new PplRelToSqlNodeConverter(MysqlSqlDialect.DEFAULT);
8686

8787
@Getter(lazy = true)
@@ -319,7 +319,7 @@ private RelNode validate(RelNode relNode, CalcitePlanContext context) {
319319
RelNode sqlRelNode = relNode.accept(new PplRelToSqlRelShuttle(context.rexBuilder, true));
320320

321321
// Convert RelNode to SqlNode for validation
322-
SqlImplementor.Result result = converter.visitRoot(sqlRelNode);
322+
SqlImplementor.Result result = rel2sql.visitRoot(sqlRelNode);
323323
SqlNode root = result.asStatement();
324324

325325
// Rewrite SqlNode to remove database qualifiers
@@ -363,7 +363,7 @@ public SqlNode visit(SqlIdentifier id) {
363363
validator,
364364
catalogReader,
365365
cluster,
366-
StandardConvertletTable.INSTANCE,
366+
PplConvertletTable.INSTANCE,
367367
SqlToRelConverter.config());
368368
RelRoot validatedRelRoot = sql2rel.convertQuery(rewritten, false, true);
369369
return validatedRelRoot.rel.accept(new PplRelToSqlRelShuttle(context.rexBuilder, false));

0 commit comments

Comments
 (0)