Skip to content

Commit 70aabfe

Browse files
committed
feat(api): Add post-analysis datetime UDT normalization 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/TIMESTAMP types via call.clone(). 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. No changes to core/ module — UDF definitions and implementors are untouched. 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 70aabfe

7 files changed

Lines changed: 393 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: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
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 org.apache.calcite.rel.RelHomogeneousShuttle;
13+
import org.apache.calcite.rel.RelNode;
14+
import org.apache.calcite.rel.logical.LogicalProject;
15+
import org.apache.calcite.rel.type.RelDataType;
16+
import org.apache.calcite.rel.type.RelDataTypeFactory;
17+
import org.apache.calcite.rel.type.RelDataTypeField;
18+
import org.apache.calcite.rex.RexBuilder;
19+
import org.apache.calcite.rex.RexNode;
20+
import org.apache.calcite.sql.type.SqlTypeName;
21+
22+
/** Wraps the root output with CAST(datetime → VARCHAR) for PPL wire-format compatibility. */
23+
class DatetimeOutputCastRule extends RelHomogeneousShuttle {
24+
25+
static final DatetimeOutputCastRule INSTANCE = new DatetimeOutputCastRule();
26+
27+
@Override
28+
public RelNode visit(RelNode other) {
29+
List<RelDataTypeField> fields = other.getRowType().getFieldList();
30+
if (fields.stream().noneMatch(f -> isDatetimeType(f.getType().getSqlTypeName()))) {
31+
return other;
32+
}
33+
34+
RexBuilder rexBuilder = other.getCluster().getRexBuilder();
35+
List<RexNode> projects = new ArrayList<>(fields.size());
36+
List<String> names = new ArrayList<>(fields.size());
37+
38+
for (RelDataTypeField field : fields) {
39+
RexNode ref = rexBuilder.makeInputRef(other, field.getIndex());
40+
if (isDatetimeType(field.getType().getSqlTypeName())) {
41+
projects.add(castToVarchar(rexBuilder, ref, field.getType()));
42+
} else {
43+
projects.add(ref);
44+
}
45+
names.add(field.getName());
46+
}
47+
return LogicalProject.create(other, List.of(), projects, names);
48+
}
49+
50+
private static RexNode castToVarchar(RexBuilder rexBuilder, RexNode expr, RelDataType fieldType) {
51+
RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
52+
RelDataType varcharType =
53+
typeFactory.createTypeWithNullability(
54+
typeFactory.createSqlType(SqlTypeName.VARCHAR), fieldType.isNullable());
55+
return rexBuilder.makeCast(varcharType, expr);
56+
}
57+
}
Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
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 org.apache.calcite.rel.RelHomogeneousShuttle;
10+
import org.apache.calcite.rel.RelNode;
11+
import org.apache.calcite.rel.type.RelDataType;
12+
import org.apache.calcite.rel.type.RelDataTypeFactory;
13+
import org.apache.calcite.rex.RexBuilder;
14+
import org.apache.calcite.rex.RexCall;
15+
import org.apache.calcite.rex.RexNode;
16+
import org.apache.calcite.rex.RexShuttle;
17+
import org.opensearch.sql.api.spec.datetime.DatetimeExtension.UdtMapping;
18+
19+
/**
20+
* Temporary patch that rewrites datetime UDT return types on RexCall nodes to standard Calcite
21+
* types.
22+
*/
23+
class DatetimeUdtNormalizeRule extends RelHomogeneousShuttle {
24+
25+
static final DatetimeUdtNormalizeRule INSTANCE = new DatetimeUdtNormalizeRule();
26+
27+
@Override
28+
public RelNode visit(RelNode other) {
29+
RelNode visited = super.visit(other);
30+
RexBuilder rexBuilder = visited.getCluster().getRexBuilder();
31+
RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
32+
return visited.accept(
33+
new RexShuttle() {
34+
@Override
35+
public RexNode visitCall(RexCall call) {
36+
call = (RexCall) super.visitCall(call);
37+
Optional<UdtMapping> mapping = UdtMapping.fromUdtType(call.getType());
38+
if (mapping.isEmpty()) {
39+
return call;
40+
}
41+
42+
RelDataType stdType =
43+
typeFactory.createTypeWithNullability(
44+
typeFactory.createSqlType(mapping.get().getStdType()),
45+
call.getType().isNullable());
46+
return call.clone(stdType, call.getOperands());
47+
}
48+
});
49+
}
50+
}

0 commit comments

Comments
 (0)