Skip to content

Commit 4a420d3

Browse files
yuancuyuancu
committed
Add special handling for datetime coercion
Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 0629a03 commit 4a420d3

5 files changed

Lines changed: 124 additions & 13 deletions

File tree

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

Lines changed: 46 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,25 @@
55

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

8+
import static org.opensearch.sql.calcite.validate.ValidationUtils.createUDTWithAttributes;
9+
810
import java.util.List;
911
import java.util.Map;
1012
import java.util.Set;
1113
import java.util.stream.IntStream;
1214
import org.apache.calcite.rel.type.RelDataType;
1315
import org.apache.calcite.rel.type.RelDataTypeFactory;
1416
import org.apache.calcite.sql.SqlCallBinding;
17+
import org.apache.calcite.sql.SqlNode;
1518
import org.apache.calcite.sql.type.SqlTypeFamily;
19+
import org.apache.calcite.sql.type.SqlTypeMappingRule;
20+
import org.apache.calcite.sql.type.SqlTypeName;
21+
import org.apache.calcite.sql.type.SqlTypeUtil;
1622
import org.apache.calcite.sql.validate.SqlValidator;
23+
import org.apache.calcite.sql.validate.SqlValidatorScope;
1724
import org.apache.calcite.sql.validate.implicit.TypeCoercionImpl;
25+
import org.checkerframework.checker.nullness.qual.Nullable;
26+
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
1827

1928
/**
2029
* Custom type coercion implementation for PPL that extends Calcite's default type coercion with
@@ -30,14 +39,14 @@ public class PplTypeCoercion extends TypeCoercionImpl {
3039

3140
static {
3241
// Initialize the blacklist for coercions that are not allowed in PPL.
33-
BLACKLISTED_COERCIONS =
34-
Map.of(
35-
SqlTypeFamily.CHARACTER,
36-
Set.of(SqlTypeFamily.NUMERIC),
37-
SqlTypeFamily.STRING,
38-
Set.of(SqlTypeFamily.NUMERIC),
39-
SqlTypeFamily.NUMERIC,
40-
Set.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.STRING));
42+
BLACKLISTED_COERCIONS = Map.of();
43+
// Map.of(
44+
// SqlTypeFamily.CHARACTER,
45+
// Set.of(SqlTypeFamily.NUMERIC),
46+
// SqlTypeFamily.STRING,
47+
// Set.of(SqlTypeFamily.NUMERIC),
48+
// SqlTypeFamily.NUMERIC,
49+
// Set.of(SqlTypeFamily.CHARACTER, SqlTypeFamily.STRING));
4150
}
4251

4352
public PplTypeCoercion(RelDataTypeFactory typeFactory, SqlValidator validator) {
@@ -73,4 +82,33 @@ private boolean isBlacklistedCoercion(RelDataType operandType, SqlTypeFamily exp
7382
}
7483
return false;
7584
}
85+
86+
@Override
87+
public @Nullable RelDataType implicitCast(RelDataType in, SqlTypeFamily expected) {
88+
RelDataType casted = super.implicitCast(in, expected);
89+
return switch (casted.getSqlTypeName()) {
90+
case SqlTypeName.DATE ->
91+
createUDTWithAttributes(factory, in, OpenSearchTypeFactory.ExprUDT.EXPR_DATE);
92+
case SqlTypeName.TIME ->
93+
createUDTWithAttributes(factory, in, OpenSearchTypeFactory.ExprUDT.EXPR_TIME);
94+
case SqlTypeName.TIMESTAMP ->
95+
createUDTWithAttributes(factory, in, OpenSearchTypeFactory.ExprUDT.EXPR_TIMESTAMP);
96+
default -> casted;
97+
};
98+
}
99+
100+
/**
101+
* Override super implementation to add special handling for user-defined types (UDTs). Otherwise,
102+
* UDTs will be regarded as character types, invalidating string->datetime casts.
103+
*/
104+
@Override
105+
protected boolean needToCast(
106+
SqlValidatorScope scope, SqlNode node, RelDataType toType, SqlTypeMappingRule mappingRule) {
107+
boolean need = super.needToCast(scope, node, toType, mappingRule);
108+
RelDataType fromType = validator.deriveType(scope, node);
109+
if (OpenSearchTypeFactory.isUserDefinedType(toType) && SqlTypeUtil.isCharacter(fromType)) {
110+
need = true;
111+
}
112+
return need;
113+
}
76114
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.validate;
7+
8+
import org.apache.calcite.sql.type.SqlTypeCoercionRule;
9+
10+
public class PplTypeCoercionRules extends SqlTypeCoercionRule {
11+
}

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55

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

8+
import static org.opensearch.sql.calcite.validate.ValidationUtils.createUDTWithAttributes;
9+
810
import java.util.List;
911
import java.util.function.Function;
1012
import org.apache.calcite.rel.type.RelDataType;
@@ -94,10 +96,10 @@ private RelDataType sqlTypeToUserDefinedType(RelDataType type) {
9496
t -> {
9597
OpenSearchTypeFactory typeFactory = (OpenSearchTypeFactory) this.getTypeFactory();
9698
return switch (t.getSqlTypeName()) {
97-
case TIMESTAMP -> typeFactory.createUDT(ExprUDT.EXPR_TIMESTAMP, t.isNullable());
98-
case TIME -> typeFactory.createUDT(ExprUDT.EXPR_TIME, t.isNullable());
99-
case DATE -> typeFactory.createUDT(ExprUDT.EXPR_DATE, t.isNullable());
100-
case BINARY -> typeFactory.createUDT(ExprUDT.EXPR_BINARY, t.isNullable());
99+
case TIMESTAMP -> createUDTWithAttributes(typeFactory, t, ExprUDT.EXPR_TIMESTAMP);
100+
case TIME -> createUDTWithAttributes(typeFactory, t, ExprUDT.EXPR_TIME);
101+
case DATE -> createUDTWithAttributes(typeFactory, t, ExprUDT.EXPR_DATE);
102+
case BINARY -> createUDTWithAttributes(typeFactory, t, ExprUDT.EXPR_BINARY);
101103
default -> t;
102104
};
103105
});

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

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.validate;
7+
8+
import static org.apache.calcite.sql.type.NonNullableAccessors.getCollation;
9+
10+
import java.nio.charset.Charset;
11+
import lombok.experimental.UtilityClass;
12+
import org.apache.calcite.rel.type.RelDataType;
13+
import org.apache.calcite.rel.type.RelDataTypeFactory;
14+
import org.apache.calcite.sql.SqlCollation;
15+
import org.apache.calcite.sql.type.SqlTypeUtil;
16+
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
17+
18+
@UtilityClass
19+
public class ValidationUtils {
20+
/**
21+
* Sync the nullability, collation, etc. to the target type. Copied from {@link
22+
* org.apache.calcite.sql.validate.implicit.AbstractTypeCoercion}
23+
*/
24+
public static RelDataType syncAttributes(
25+
RelDataTypeFactory factory, RelDataType fromType, RelDataType toType) {
26+
RelDataType syncedType = toType;
27+
if (fromType != null) {
28+
syncedType = factory.createTypeWithNullability(syncedType, fromType.isNullable());
29+
if (SqlTypeUtil.inCharOrBinaryFamilies(fromType)
30+
&& SqlTypeUtil.inCharOrBinaryFamilies(toType)) {
31+
Charset charset = fromType.getCharset();
32+
if (charset != null && SqlTypeUtil.inCharFamily(syncedType)) {
33+
SqlCollation collation = getCollation(fromType);
34+
syncedType = factory.createTypeWithCharsetAndCollation(syncedType, charset, collation);
35+
}
36+
}
37+
}
38+
return syncedType;
39+
}
40+
41+
/**
42+
* Creates a user-defined type with attributes (nullability, charset, collation) copied from
43+
* another type.
44+
*
45+
* @param factory the type factory used to create the UDT
46+
* @param fromType the source type to copy attributes from (nullability, charset, collation)
47+
* @param userDefinedType the user-defined type to create
48+
* @return a new RelDataType representing the UDT with attributes from fromType
49+
*/
50+
public static RelDataType createUDTWithAttributes(
51+
RelDataTypeFactory factory,
52+
RelDataType fromType,
53+
OpenSearchTypeFactory.ExprUDT userDefinedType) {
54+
if (!(factory instanceof OpenSearchTypeFactory typeFactory)) {
55+
throw new IllegalArgumentException("factory must be an instance of OpenSearchTypeFactory");
56+
}
57+
RelDataType type = typeFactory.createUDT(userDefinedType);
58+
return syncAttributes(typeFactory, fromType, type);
59+
}
60+
}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ public class OpenSearchExecutionEngine implements ExecutionEngine {
7777
private final PlanSerializer planSerializer;
7878

7979
static {
80-
CalcitePlanContext.setOperatorTableProvider(OperatorTable::getChainedOperatorTable);
80+
CalcitePlanContext.setOperatorTableProvider(OperatorTable::getChainedOperatorTable);
8181
}
8282

8383
public OpenSearchExecutionEngine(

0 commit comments

Comments
 (0)