Skip to content

Commit da7fd7c

Browse files
coderabbitai[bot]coderabbitai[bot]
authored
📝 Add docstrings to issues/4636
Docstrings generation was requested by @yuancu. * #4892 (comment) The following files were modified: * `core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java` * `core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java` * `core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java` * `core/src/main/java/org/opensearch/sql/calcite/ExtendedRexBuilder.java` * `core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java` * `core/src/main/java/org/opensearch/sql/calcite/utils/OpenSearchTypeFactory.java` * `core/src/main/java/org/opensearch/sql/calcite/utils/OpenSearchTypeUtil.java` * `core/src/main/java/org/opensearch/sql/calcite/utils/binning/BinnableField.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/OpenSearchSparkSqlDialect.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/PplConvertletTable.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/PplTypeCoercion.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/PplTypeCoercionRule.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/PplValidator.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/SqlOperatorTableProvider.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/ValidationUtils.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/converters/PplRelToSqlNodeConverter.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/converters/PplSqlToRelConverter.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/shuttles/PplRelToSqlRelShuttle.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/shuttles/SkipRelValidationShuttle.java` * `core/src/main/java/org/opensearch/sql/calcite/validate/shuttles/SqlRewriteShuttle.java` * `core/src/main/java/org/opensearch/sql/executor/QueryService.java` * `core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ArrayFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ExistsFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/FilterFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ForallFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/LambdaUtils.java` * `core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/MVAppendFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/MapAppendFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/MapRemoveFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/ReduceFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/TransformFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java` * `core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java` * `core/src/main/java/org/opensearch/sql/expression/function/UDFOperandMetadata.java` * `core/src/main/java/org/opensearch/sql/expression/function/UserDefinedFunctionBuilder.java` * `core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonAppendFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonDeleteFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonExtendFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonExtractAllFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonExtractFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonKeysFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/jsonUDF/JsonSetFunctionImpl.java` * `core/src/main/java/org/opensearch/sql/expression/function/udf/RelevanceQueryFunction.java` * `core/src/main/java/org/opensearch/sql/expression/function/udf/SpanFunction.java` * `core/src/main/java/org/opensearch/sql/expression/function/udf/binning/WidthBucketFunction.java` * `core/src/main/java/org/opensearch/sql/expression/function/udf/math/DivideFunction.java` * `core/src/main/java/org/opensearch/sql/expression/function/udf/math/ModFunction.java` * `core/src/main/java/org/opensearch/sql/expression/function/udf/math/ScalarMaxFunction.java` * `core/src/main/java/org/opensearch/sql/expression/function/udf/math/ScalarMinFunction.java` * `core/src/test/java/org/opensearch/sql/expression/function/AggFunctionTestBase.java` * `integ-test/src/test/java/org/opensearch/sql/calcite/clickbench/PPLClickBenchIT.java` * `integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteArrayFunctionIT.java` * `integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java` * `integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteMultisearchCommandIT.java` * `integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLAppendCommandIT.java` * `integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLAppendcolIT.java` * `integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLBasicIT.java` * `integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLExplainIT.java` * `integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLPatternsIT.java` * `integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteWhereCommandIT.java` * `integ-test/src/test/java/org/opensearch/sql/ppl/DateTimeFunctionIT.java` * `integ-test/src/test/java/org/opensearch/sql/ppl/ExplainIT.java`
1 parent cbcdbd6 commit da7fd7c

63 files changed

Lines changed: 3985 additions & 1185 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java‎

Lines changed: 87 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY;
99

1010
import java.sql.Connection;
11+
import java.sql.SQLException;
1112
import java.util.ArrayList;
1213
import java.util.HashMap;
1314
import java.util.List;
@@ -17,13 +18,21 @@
1718
import java.util.function.BiFunction;
1819
import lombok.Getter;
1920
import lombok.Setter;
21+
import org.apache.calcite.config.NullCollation;
2022
import org.apache.calcite.rex.RexCorrelVariable;
2123
import org.apache.calcite.rex.RexLambdaRef;
2224
import org.apache.calcite.rex.RexNode;
25+
import org.apache.calcite.server.CalciteServerStatement;
26+
import org.apache.calcite.sql.validate.SqlValidator;
2327
import org.apache.calcite.tools.FrameworkConfig;
2428
import org.apache.calcite.tools.RelBuilder;
2529
import org.opensearch.sql.ast.expression.UnresolvedExpression;
2630
import org.opensearch.sql.calcite.utils.CalciteToolsHelper;
31+
import org.opensearch.sql.calcite.validate.OpenSearchSparkSqlDialect;
32+
import org.opensearch.sql.calcite.validate.PplTypeCoercion;
33+
import org.opensearch.sql.calcite.validate.PplTypeCoercionRule;
34+
import org.opensearch.sql.calcite.validate.PplValidator;
35+
import org.opensearch.sql.calcite.validate.SqlOperatorTableProvider;
2736
import org.opensearch.sql.common.setting.Settings;
2837
import org.opensearch.sql.executor.QueryType;
2938
import org.opensearch.sql.expression.function.FunctionProperties;
@@ -72,6 +81,25 @@ public class CalcitePlanContext {
7281
/** Whether we're currently inside a lambda context. */
7382
@Getter @Setter private boolean inLambdaContext = false;
7483

84+
/**
85+
* -- SETTER -- Sets the SQL operator table provider. This must be called during initialization by
86+
* the opensearch module.
87+
*
88+
* @param provider the provider to use for obtaining operator tables
89+
*/
90+
@Setter private static SqlOperatorTableProvider operatorTableProvider;
91+
92+
/** Cached SqlValidator instance (lazy initialized). */
93+
private SqlValidator validator;
94+
95+
/**
96+
* Initialize a new CalcitePlanContext and set up Calcite connection, builders, function
97+
* properties, and lambda/capture state.
98+
*
99+
* @param config the Calcite FrameworkConfig to use for planning
100+
* @param sysLimit system limits governing planning behaviour
101+
* @param queryType the QueryType for this planning context
102+
*/
75103
private CalcitePlanContext(FrameworkConfig config, SysLimit sysLimit, QueryType queryType) {
76104
this.config = config;
77105
this.sysLimit = sysLimit;
@@ -85,8 +113,11 @@ private CalcitePlanContext(FrameworkConfig config, SysLimit sysLimit, QueryType
85113
}
86114

87115
/**
88-
* Private constructor for creating a context that shares relBuilder with parent. Used by clone()
89-
* to create lambda contexts that can resolve fields from the parent context.
116+
* Creates a child CalcitePlanContext that shares relational builders with the given parent while isolating lambda-specific state.
117+
*
118+
* <p>The child context reuses the parent's relBuilder and rexBuilder so it can resolve fields against the parent, but it initializes independent maps/lists for lambda references and captured variables and marks the context as a lambda context.</p>
119+
*
120+
* @param parent the parent context whose builders and configuration are reused
90121
*/
91122
private CalcitePlanContext(CalcitePlanContext parent) {
92123
this.config = parent.config;
@@ -101,6 +132,59 @@ private CalcitePlanContext(CalcitePlanContext parent) {
101132
this.inLambdaContext = true; // Mark that we're inside a lambda
102133
}
103134

135+
/**
136+
* Provides the lazily initialized SqlValidator configured for PPL within this context.
137+
*
138+
* @return the SqlValidator configured with PPL type coercion rules, Spark SQL conformance,
139+
* Spark NULL collation, and identifier expansion
140+
* @throws IllegalStateException if the global SqlOperatorTableProvider has not been set
141+
* @throws RuntimeException if creating or unwrapping the underlying Calcite statement fails
142+
*/
143+
public SqlValidator getValidator() {
144+
if (validator == null) {
145+
final CalciteServerStatement statement;
146+
try {
147+
statement = connection.createStatement().unwrap(CalciteServerStatement.class);
148+
} catch (SQLException e) {
149+
throw new RuntimeException(e);
150+
}
151+
if (operatorTableProvider == null) {
152+
throw new IllegalStateException(
153+
"SqlOperatorTableProvider must be set before creating CalcitePlanContext");
154+
}
155+
SqlValidator.Config validatorConfig =
156+
SqlValidator.Config.DEFAULT
157+
.withTypeCoercionRules(PplTypeCoercionRule.instance())
158+
.withTypeCoercionFactory(PplTypeCoercion::create)
159+
// Use lenient conformance for PPL compatibility
160+
.withConformance(OpenSearchSparkSqlDialect.DEFAULT.getConformance())
161+
// Use Spark SQL's NULL collation (NULLs sorted LOW/FIRST)
162+
.withDefaultNullCollation(NullCollation.LOW)
163+
// This ensures that coerced arguments are replaced with cast version in sql select
164+
// list because coercion is performed during select list expansion during sql
165+
// validation. Affects 4356.yml
166+
// See SqlValidatorImpl#validateSelectList and AggConverter#translateAgg
167+
.withIdentifierExpansion(true);
168+
validator =
169+
PplValidator.create(
170+
statement,
171+
config,
172+
operatorTableProvider.getOperatorTable(),
173+
TYPE_FACTORY,
174+
validatorConfig);
175+
}
176+
return validator;
177+
}
178+
179+
/**
180+
* Temporarily marks the context as resolving a join condition and applies a transformation to the given unresolved expression.
181+
*
182+
* <p>The context flag indicating join-condition resolution is set to true before invoking the transform and reset to false afterwards.
183+
*
184+
* @param expr the unresolved expression representing the join condition
185+
* @param transformFunction a function that converts the unresolved expression into a {@code RexNode} using this context
186+
* @return the {@code RexNode} produced by applying {@code transformFunction} to {@code expr}
187+
*/
104188
public RexNode resolveJoinCondition(
105189
UnresolvedExpression expr,
106190
BiFunction<UnresolvedExpression, CalcitePlanContext, RexNode> transformFunction) {
@@ -206,4 +290,4 @@ public RexLambdaRef captureVariable(RexNode fieldRef, String fieldName) {
206290

207291
return lambdaRef;
208292
}
209-
}
293+
}

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java‎

Lines changed: 115 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,12 @@
4343
import java.util.stream.IntStream;
4444
import java.util.stream.Stream;
4545
import lombok.AllArgsConstructor;
46+
import lombok.NonNull;
4647
import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
4748
import org.apache.calcite.plan.RelOptTable;
4849
import org.apache.calcite.plan.ViewExpanders;
50+
import org.apache.calcite.rel.RelCollation;
51+
import org.apache.calcite.rel.RelFieldCollation;
4952
import org.apache.calcite.rel.RelNode;
5053
import org.apache.calcite.rel.core.Aggregate;
5154
import org.apache.calcite.rel.core.JoinRelType;
@@ -55,10 +58,14 @@
5558
import org.apache.calcite.rel.type.RelDataTypeField;
5659
import org.apache.calcite.rex.RexCall;
5760
import org.apache.calcite.rex.RexCorrelVariable;
61+
import org.apache.calcite.rex.RexFieldCollation;
5862
import org.apache.calcite.rex.RexInputRef;
5963
import org.apache.calcite.rex.RexLiteral;
6064
import org.apache.calcite.rex.RexNode;
65+
import org.apache.calcite.rex.RexOver;
66+
import org.apache.calcite.rex.RexShuttle;
6167
import org.apache.calcite.rex.RexVisitorImpl;
68+
import org.apache.calcite.rex.RexWindow;
6269
import org.apache.calcite.rex.RexWindowBounds;
6370
import org.apache.calcite.sql.SqlKind;
6471
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
@@ -1678,6 +1685,17 @@ private void validateFillNullTypeCompatibility(
16781685
}
16791686
}
16801687

1688+
/**
1689+
* Builds a Calcite relational plan that implements the streaming window operation defined by the given AST node.
1690+
*
1691+
* The produced plan handles reset-before/reset-after semantics, global sliding windows, grouped windows, and
1692+
* per-window function expressions; it also embeds existing plan collations into window OVER clauses and
1693+
* adds any required helper columns (sequence/segment/reset flags) used by the stream-window logic.
1694+
*
1695+
* @param node the StreamWindow AST node describing window type, grouping, window size, reset options and functions
1696+
* @param context the planning context containing the RelBuilder, visitors, and planner state
1697+
* @return a RelNode representing the compiled stream window operation ready for integration into the surrounding plan
1698+
*/
16811699
@Override
16821700
public RelNode visitStreamWindow(StreamWindow node, CalcitePlanContext context) {
16831701
visitChildren(node, context);
@@ -1740,6 +1758,7 @@ public RelNode visitStreamWindow(StreamWindow node, CalcitePlanContext context)
17401758
// Default: first get rawExpr
17411759
List<RexNode> overExpressions =
17421760
node.getWindowFunctionList().stream().map(w -> rexVisitor.analyze(w, context)).toList();
1761+
overExpressions = embedExistingCollationsIntoOver(overExpressions, context);
17431762

17441763
if (hasGroup) {
17451764
// only build sequence when there is by condition
@@ -1781,6 +1800,101 @@ public RelNode visitStreamWindow(StreamWindow node, CalcitePlanContext context)
17811800
return context.relBuilder.peek();
17821801
}
17831802

1803+
/**
1804+
* Embed the current plan's collation into any RexOver window expressions.
1805+
*
1806+
* <p>Window frames such as ROWS n PRECEDING require an ORDER BY to define row order; this method
1807+
* adds the RelBuilder's existing collation as the OVER clause's ORDER BY for each encountered
1808+
* RexOver.
1809+
*
1810+
* @param overExpressions window function expressions that may contain nested {@link RexOver}
1811+
* @param context plan context used to obtain the current RelBuilder and RexBuilder
1812+
* @return a list of expressions where each {@link RexOver} has its ORDER BY populated from the
1813+
* current plan collation (unchanged expressions are returned as-is)
1814+
*/
1815+
private List<RexNode> embedExistingCollationsIntoOver(
1816+
List<RexNode> overExpressions, CalcitePlanContext context) {
1817+
RelCollation existingCollation = context.relBuilder.peek().getTraitSet().getCollation();
1818+
List<@NonNull RelFieldCollation> relCollations =
1819+
existingCollation == null ? List.of() : existingCollation.getFieldCollations();
1820+
ImmutableList<@NonNull RexFieldCollation> rexCollations =
1821+
relCollations.stream()
1822+
.map(f -> relCollationToRexCollation(f, context.relBuilder))
1823+
.collect(ImmutableList.toImmutableList());
1824+
return overExpressions.stream()
1825+
.map(
1826+
n ->
1827+
n.accept(
1828+
new RexShuttle() {
1829+
/**
1830+
* Rewrite the given windowed aggregation to attach the visitor's collected collations.
1831+
*
1832+
* @param over the original RexOver node to be rewritten
1833+
* @return a RexNode representing the same windowed aggregation with the visitor's
1834+
* accumulated RexFieldCollation ordering embedded into its OVER specification
1835+
*/
1836+
@Override
1837+
public RexNode visitOver(RexOver over) {
1838+
RexWindow window = over.getWindow();
1839+
return context.rexBuilder.makeOver(
1840+
over.getType(),
1841+
over.getAggOperator(),
1842+
over.getOperands(),
1843+
window.partitionKeys,
1844+
rexCollations,
1845+
window.getLowerBound(),
1846+
window.getUpperBound(),
1847+
window.isRows(),
1848+
true,
1849+
false,
1850+
over.isDistinct(),
1851+
over.ignoreNulls());
1852+
}
1853+
}))
1854+
.toList();
1855+
}
1856+
1857+
/**
1858+
* Create a RexFieldCollation that mirrors the given RelFieldCollation by referencing
1859+
* the corresponding field through the provided RelBuilder.
1860+
*
1861+
* @param relCollation the source RelFieldCollation whose field index, direction,
1862+
* and null ordering will be translated
1863+
* @param builder RelBuilder used to produce a RexNode referencing the field
1864+
* @return a RexFieldCollation with equivalent sort direction and null-ordering flags
1865+
*/
1866+
private static RexFieldCollation relCollationToRexCollation(
1867+
RelFieldCollation relCollation, RelBuilder builder) {
1868+
RexNode fieldRef = builder.field(relCollation.getFieldIndex());
1869+
1870+
// Convert direction flags to SqlKind set
1871+
Set<SqlKind> flags = new HashSet<>();
1872+
if (relCollation.direction == RelFieldCollation.Direction.DESCENDING
1873+
|| relCollation.direction == RelFieldCollation.Direction.STRICTLY_DESCENDING) {
1874+
flags.add(SqlKind.DESCENDING);
1875+
}
1876+
if (relCollation.nullDirection == RelFieldCollation.NullDirection.FIRST) {
1877+
flags.add(SqlKind.NULLS_FIRST);
1878+
} else if (relCollation.nullDirection == RelFieldCollation.NullDirection.LAST) {
1879+
flags.add(SqlKind.NULLS_LAST);
1880+
}
1881+
1882+
return new RexFieldCollation(fieldRef, flags);
1883+
}
1884+
1885+
/**
1886+
* Wraps each window expression in a CASE that yields NULL when the provided group-not-null
1887+
* predicate is false, preserving any explicit alias on the original expression.
1888+
*
1889+
* <p>Each expression is transformed to: CASE groupNotNull WHEN TRUE THEN <original> ELSE NULL END.
1890+
*
1891+
* @param overExpressions list of window expressions to wrap; if an expression is of the form
1892+
* `expr AS alias` the alias is preserved on the wrapped expression
1893+
* @param groupNotNull predicate used as the CASE condition to determine non-null groups
1894+
* @param context plan context providing RexBuilder/RelBuilder utilities
1895+
* @return a new list of `RexNode` where each input expression has been wrapped with the CASE
1896+
* nulling logic and retains its original alias when present
1897+
*/
17841898
private List<RexNode> wrapWindowFunctionsWithGroupNotNull(
17851899
List<RexNode> overExpressions, RexNode groupNotNull, CalcitePlanContext context) {
17861900
List<RexNode> wrappedOverExprs = new ArrayList<>(overExpressions.size());
@@ -3499,4 +3613,4 @@ private RexNode createOptimizedTransliteration(
34993613
throw new RuntimeException("Failed to optimize sed expression: " + sedExpression, e);
35003614
}
35013615
}
3502-
}
3616+
}

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java‎

Lines changed: 11 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -462,6 +462,14 @@ public RexNode visitFunction(Function node, CalcitePlanContext context) {
462462
throw new IllegalArgumentException("Unsupported operator: " + node.getFuncName());
463463
}
464464

465+
/**
466+
* Translates a window function expression into a Calcite RexNode representing an OVER() window call.
467+
*
468+
* @param node the window function expression to translate
469+
* @param context the planning context used for resolving arguments, partitions, and frame
470+
* @return a RexNode for the windowed function call (an OVER expression)
471+
* @throws UnsupportedOperationException if the function name is not a supported window function
472+
*/
465473
@Override
466474
public RexNode visitWindowFunction(WindowFunction node, CalcitePlanContext context) {
467475
Function windowFunction = (Function) node.getFunction();
@@ -480,26 +488,8 @@ public RexNode visitWindowFunction(WindowFunction node, CalcitePlanContext conte
480488
(arguments.isEmpty() || arguments.size() == 1)
481489
? Collections.emptyList()
482490
: arguments.subList(1, arguments.size());
483-
List<RexNode> nodes =
484-
PPLFuncImpTable.INSTANCE.validateAggFunctionSignature(
485-
functionName, field, args, context.rexBuilder);
486-
return nodes != null
487-
? PlanUtils.makeOver(
488-
context,
489-
functionName,
490-
nodes.getFirst(),
491-
nodes.size() <= 1 ? Collections.emptyList() : nodes.subList(1, nodes.size()),
492-
partitions,
493-
List.of(),
494-
node.getWindowFrame())
495-
: PlanUtils.makeOver(
496-
context,
497-
functionName,
498-
field,
499-
args,
500-
partitions,
501-
List.of(),
502-
node.getWindowFrame());
491+
return PlanUtils.makeOver(
492+
context, functionName, field, args, partitions, List.of(), node.getWindowFrame());
503493
})
504494
.orElseThrow(
505495
() ->
@@ -680,4 +670,4 @@ public RexNode visitUnresolvedArgument(UnresolvedArgument node, CalcitePlanConte
680670
context.rexBuilder.makeLiteral(node.getArgName()),
681671
value);
682672
}
683-
}
673+
}

0 commit comments

Comments
 (0)