Skip to content

Commit 01fba06

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/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 01fba06

7 files changed

Lines changed: 402 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: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
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.opensearch.sql.api.spec.datetime.DatetimeExtension.UdtMapping;
20+
21+
/**
22+
* Temporary patch that rewrites datetime UDT return types on RexCall nodes to standard Calcite
23+
* types.
24+
*/
25+
@NoArgsConstructor(access = AccessLevel.PRIVATE)
26+
class DatetimeUdtNormalizeRule extends RelHomogeneousShuttle {
27+
28+
static final DatetimeUdtNormalizeRule INSTANCE = new DatetimeUdtNormalizeRule();
29+
30+
@Override
31+
public RelNode visit(RelNode other) {
32+
RelNode visited = super.visit(other);
33+
RexBuilder rexBuilder = visited.getCluster().getRexBuilder();
34+
RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
35+
return visited.accept(
36+
new RexShuttle() {
37+
@Override
38+
public RexNode visitCall(RexCall call) {
39+
call = (RexCall) super.visitCall(call);
40+
Optional<UdtMapping> mapping = UdtMapping.fromUdtType(call.getType());
41+
if (mapping.isEmpty()) {
42+
return call;
43+
}
44+
45+
// Normalize UDT return type to standard Calcite DATE/TIME/TIMESTAMP
46+
RelDataType stdType =
47+
typeFactory.createTypeWithNullability(
48+
typeFactory.createSqlType(mapping.get().getStdType()),
49+
call.getType().isNullable());
50+
return call.clone(stdType, call.getOperands());
51+
}
52+
});
53+
}
54+
}

0 commit comments

Comments
 (0)