Skip to content

Commit cc90d7f

Browse files
committed
feat(api): Normalize datetime types for unified query API
Add postAnalysisRules (List<RelShuttle>) to LanguageSpec.LanguageExtension and register DatetimeExtension in UnifiedPplSpec with two rules: 1. DatetimeUdtNormalizeRule rewrites datetime UDT return types (EXPR_DATE/TIME/TIMESTAMP) on RexCall nodes to standard Calcite DATE/TIME(9)/TIMESTAMP(9) types via call.clone(). Precision is derived from the type system (OpenSearchTypeSystem.getMaxPrecision). 2. DatetimeOutputCastRule adds a final LogicalProject that casts standard datetime output columns to VARCHAR, aligning with PPL's wire-format contract (ISO string representation). Both rules run as postAnalysisRules after the planning strategy produces the RelNode, applied uniformly to both SQL and PPL paths. Also bumps OpenSearchTypeSystem max datetime precision from 3 to 9 (nanosecond) for TIME and TIMESTAMP types. No changes to UDF definitions or implementors in core/ — the mismatch between rewritten signatures and UDF implementations is a known limitation addressed separately. Signed-off-by: Chen Dai <daichen@amazon.com>
1 parent d3bdca8 commit cc90d7f

8 files changed

Lines changed: 445 additions & 9 deletions

File tree

api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,15 @@ public UnifiedQueryPlanner(UnifiedQueryContext context) {
6060
*/
6161
public RelNode plan(String query) {
6262
try {
63-
return context.measure(ANALYZE, () -> strategy.plan(query));
63+
return context.measure(
64+
ANALYZE,
65+
() -> {
66+
RelNode plan = strategy.plan(query);
67+
for (var shuttle : context.getLangSpec().postAnalysisRules()) {
68+
plan = plan.accept(shuttle);
69+
}
70+
return plan;
71+
});
6472
} catch (SyntaxCheckException | UnsupportedOperationException e) {
6573
throw e;
6674
} catch (Exception e) {

api/src/main/java/org/opensearch/sql/api/spec/LanguageSpec.java

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import java.util.ArrayList;
99
import java.util.List;
10+
import org.apache.calcite.rel.RelShuttle;
1011
import org.apache.calcite.sql.SqlNode;
1112
import org.apache.calcite.sql.SqlOperatorTable;
1213
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
@@ -17,8 +18,8 @@
1718

1819
/**
1920
* Language specification defining the dialect the engine accepts. Provides parser configuration,
20-
* validator configuration, and composable {@link LanguageExtension}s that contribute operators and
21-
* post-parse rewrite rules.
21+
* validator configuration, and composable {@link LanguageExtension}s that contribute operators,
22+
* post-parse rewrite rules, and post-analysis rewrite rules.
2223
*
2324
* <p>Implementations define a complete language surface — for example, {@link UnifiedSqlSpec}
2425
* provides ANSI and extended SQL modes. A future PPL spec would implement this same interface once
@@ -27,8 +28,9 @@
2728
public interface LanguageSpec {
2829

2930
/**
30-
* A composable language extension that contributes operators and post-parse rewrite rules. All
31-
* methods have defaults so extensions only override what they need.
31+
* A composable language extension that contributes operators, post-parse rewrite rules, and
32+
* post-analysis rewrite rules. All methods have defaults so extensions only override what they
33+
* need.
3234
*/
3335
interface LanguageExtension {
3436

@@ -47,6 +49,14 @@ default SqlOperatorTable operators() {
4749
default List<SqlVisitor<SqlNode>> postParseRules() {
4850
return List.of();
4951
}
52+
53+
/**
54+
* RelNode rewrite rules applied after analysis and before execution. Each rule transforms the
55+
* logical plan tree. Rules within a single extension are applied in list order.
56+
*/
57+
default List<RelShuttle> postAnalysisRules() {
58+
return List.of();
59+
}
5060
}
5161

5262
/**
@@ -62,9 +72,9 @@ default List<SqlVisitor<SqlNode>> postParseRules() {
6272
SqlValidator.Config validatorConfig();
6373

6474
/**
65-
* Language extensions registered with this spec. Each extension contributes operators and
66-
* post-parse rewrite rules that are composed by {@link #operatorTable()} and {@link
67-
* #postParseRules()}.
75+
* Language extensions registered with this spec. Each extension contributes operators, post-parse
76+
* rewrite rules, and post-analysis rewrite rules composed by {@link #operatorTable()}, {@link
77+
* #postParseRules()}, and {@link #postAnalysisRules()}.
6878
*/
6979
List<LanguageExtension> extensions();
7080

@@ -86,4 +96,12 @@ default SqlOperatorTable operatorTable() {
8696
default List<SqlVisitor<SqlNode>> postParseRules() {
8797
return extensions().stream().flatMap(ext -> ext.postParseRules().stream()).toList();
8898
}
99+
100+
/**
101+
* All post-analysis RelNode rewrite rules from registered extensions, flattened in registration
102+
* order. Applied to the logical plan after analysis and before execution.
103+
*/
104+
default List<RelShuttle> postAnalysisRules() {
105+
return extensions().stream().flatMap(ext -> ext.postAnalysisRules().stream()).toList();
106+
}
89107
}

api/src/main/java/org/opensearch/sql/api/spec/UnifiedPplSpec.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import lombok.NoArgsConstructor;
1111
import org.apache.calcite.sql.parser.SqlParser;
1212
import org.apache.calcite.sql.validate.SqlValidator;
13+
import org.opensearch.sql.api.spec.datetime.DatetimeExtension;
1314

1415
/**
1516
* PPL language specification.
@@ -37,6 +38,6 @@ public SqlValidator.Config validatorConfig() {
3738

3839
@Override
3940
public List<LanguageExtension> extensions() {
40-
return List.of();
41+
return List.of(new DatetimeExtension());
4142
}
4243
}
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.api.spec.datetime;
7+
8+
import java.util.Arrays;
9+
import java.util.List;
10+
import java.util.Optional;
11+
import lombok.Getter;
12+
import lombok.RequiredArgsConstructor;
13+
import org.apache.calcite.rel.RelShuttle;
14+
import org.apache.calcite.rel.type.RelDataType;
15+
import org.apache.calcite.sql.type.SqlTypeName;
16+
import org.opensearch.sql.api.spec.LanguageSpec.LanguageExtension;
17+
import org.opensearch.sql.calcite.type.AbstractExprRelDataType;
18+
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.ExprUDT;
19+
20+
/** Datetime language extension that normalizes UDT types and casts output for wire-format. */
21+
public class DatetimeExtension implements LanguageExtension {
22+
23+
@Override
24+
public List<RelShuttle> postAnalysisRules() {
25+
return List.of(DatetimeUdtNormalizeRule.INSTANCE, DatetimeOutputCastRule.INSTANCE);
26+
}
27+
28+
/** Maps datetime UDT types to their standard Calcite equivalents. */
29+
@Getter
30+
@RequiredArgsConstructor
31+
enum UdtMapping {
32+
DATE(ExprUDT.EXPR_DATE, SqlTypeName.DATE),
33+
TIME(ExprUDT.EXPR_TIME, SqlTypeName.TIME),
34+
TIMESTAMP(ExprUDT.EXPR_TIMESTAMP, SqlTypeName.TIMESTAMP);
35+
36+
private final ExprUDT udtType;
37+
private final SqlTypeName stdType;
38+
39+
/** Matches a UDT RelDataType to its mapping, or empty if not a datetime UDT. */
40+
static Optional<UdtMapping> fromUdtType(RelDataType type) {
41+
if (!(type instanceof AbstractExprRelDataType<?> e)) {
42+
return Optional.empty();
43+
}
44+
ExprUDT udt = e.getUdt();
45+
return Arrays.stream(values()).filter(u -> u.udtType == udt).findFirst();
46+
}
47+
48+
/** Returns true if the given SqlTypeName is a standard datetime type. */
49+
static boolean isDatetimeType(SqlTypeName typeName) {
50+
return Arrays.stream(values()).anyMatch(u -> u.stdType == typeName);
51+
}
52+
}
53+
}
Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.api.spec.datetime;
7+
8+
import static org.opensearch.sql.api.spec.datetime.DatetimeExtension.UdtMapping.isDatetimeType;
9+
10+
import java.util.ArrayList;
11+
import java.util.List;
12+
import lombok.AccessLevel;
13+
import lombok.NoArgsConstructor;
14+
import org.apache.calcite.rel.RelHomogeneousShuttle;
15+
import org.apache.calcite.rel.RelNode;
16+
import org.apache.calcite.rel.logical.LogicalProject;
17+
import org.apache.calcite.rel.type.RelDataType;
18+
import org.apache.calcite.rel.type.RelDataTypeFactory;
19+
import org.apache.calcite.rel.type.RelDataTypeField;
20+
import org.apache.calcite.rex.RexBuilder;
21+
import org.apache.calcite.rex.RexNode;
22+
import org.apache.calcite.sql.type.SqlTypeName;
23+
24+
/** Wraps the root output with CAST(datetime → VARCHAR) for PPL wire-format compatibility. */
25+
@NoArgsConstructor(access = AccessLevel.PRIVATE)
26+
class DatetimeOutputCastRule extends RelHomogeneousShuttle {
27+
28+
static final DatetimeOutputCastRule INSTANCE = new DatetimeOutputCastRule();
29+
30+
@Override
31+
public RelNode visit(RelNode other) {
32+
List<RelDataTypeField> fields = other.getRowType().getFieldList();
33+
if (fields.stream().noneMatch(f -> isDatetimeType(f.getType().getSqlTypeName()))) {
34+
return other;
35+
}
36+
37+
RexBuilder rexBuilder = other.getCluster().getRexBuilder();
38+
List<RexNode> projects = new ArrayList<>(fields.size());
39+
List<String> names = new ArrayList<>(fields.size());
40+
41+
// Cast datetime fields to VARCHAR for output; pass through others unchanged
42+
for (RelDataTypeField field : fields) {
43+
RexNode newField = rexBuilder.makeInputRef(other, field.getIndex());
44+
RelDataType fieldType = field.getType();
45+
if (isDatetimeType(fieldType.getSqlTypeName())) {
46+
projects.add(castToVarchar(rexBuilder, newField, fieldType));
47+
} else {
48+
projects.add(newField);
49+
}
50+
names.add(field.getName());
51+
}
52+
return LogicalProject.create(other, List.of(), projects, names);
53+
}
54+
55+
private static RexNode castToVarchar(RexBuilder rexBuilder, RexNode expr, RelDataType fieldType) {
56+
RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
57+
RelDataType varcharType =
58+
typeFactory.createTypeWithNullability(
59+
typeFactory.createSqlType(SqlTypeName.VARCHAR), fieldType.isNullable());
60+
return rexBuilder.makeCast(varcharType, expr);
61+
}
62+
}
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.api.spec.datetime;
7+
8+
import java.util.Optional;
9+
import lombok.AccessLevel;
10+
import lombok.NoArgsConstructor;
11+
import org.apache.calcite.rel.RelHomogeneousShuttle;
12+
import org.apache.calcite.rel.RelNode;
13+
import org.apache.calcite.rel.type.RelDataType;
14+
import org.apache.calcite.rel.type.RelDataTypeFactory;
15+
import org.apache.calcite.rex.RexBuilder;
16+
import org.apache.calcite.rex.RexCall;
17+
import org.apache.calcite.rex.RexNode;
18+
import org.apache.calcite.rex.RexShuttle;
19+
import org.apache.calcite.sql.type.SqlTypeName;
20+
import org.opensearch.sql.api.spec.datetime.DatetimeExtension.UdtMapping;
21+
22+
/**
23+
* Temporary patch that rewrites datetime UDT return types on RexCall nodes to standard Calcite
24+
* types.
25+
*/
26+
@NoArgsConstructor(access = AccessLevel.PRIVATE)
27+
class DatetimeUdtNormalizeRule extends RelHomogeneousShuttle {
28+
29+
static final DatetimeUdtNormalizeRule INSTANCE = new DatetimeUdtNormalizeRule();
30+
31+
@Override
32+
public RelNode visit(RelNode other) {
33+
RelNode visited = super.visit(other);
34+
RexBuilder rexBuilder = visited.getCluster().getRexBuilder();
35+
RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
36+
return visited.accept(
37+
new RexShuttle() {
38+
@Override
39+
public RexNode visitCall(RexCall call) {
40+
call = (RexCall) super.visitCall(call);
41+
Optional<UdtMapping> mapping = UdtMapping.fromUdtType(call.getType());
42+
if (mapping.isEmpty()) {
43+
return call;
44+
}
45+
46+
// Normalize UDT return type to standard Calcite DATE/TIME/TIMESTAMP
47+
UdtMapping m = mapping.get();
48+
SqlTypeName stdTypeName = m.getStdType();
49+
RelDataType baseType =
50+
stdTypeName.allowsPrec()
51+
? typeFactory.createSqlType(
52+
stdTypeName, typeFactory.getTypeSystem().getMaxPrecision(stdTypeName))
53+
: typeFactory.createSqlType(stdTypeName);
54+
RelDataType stdType =
55+
typeFactory.createTypeWithNullability(baseType, call.getType().isNullable());
56+
return call.clone(stdType, call.getOperands());
57+
}
58+
});
59+
}
60+
}

0 commit comments

Comments
 (0)