Skip to content

Commit c625fc6

Browse files
tanik98tanik98
authored
Wire PPL conditional functions with analytics-backend-datafusion (opensearch-project#21643)
Signed-off-by: Tanik Pansuriya <panbhai@amazon.com>
1 parent 2f2c1f4 commit c625fc6

4 files changed

Lines changed: 585 additions & 0 deletions

File tree

‎sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/CoalesceAdapter.java‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.be.datafusion;
10+
11+
import org.apache.calcite.plan.RelOptCluster;
12+
import org.apache.calcite.rel.type.RelDataType;
13+
import org.apache.calcite.rel.type.RelDataTypeFactory;
14+
import org.apache.calcite.rex.RexBuilder;
15+
import org.apache.calcite.rex.RexCall;
16+
import org.apache.calcite.rex.RexNode;
17+
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
18+
import org.apache.calcite.sql.type.SqlTypeName;
19+
import org.opensearch.analytics.spi.FieldStorageInfo;
20+
import org.opensearch.analytics.spi.ScalarFunctionAdapter;
21+
22+
import java.util.ArrayList;
23+
import java.util.List;
24+
25+
/**
26+
* Adapter for coalesce(field1, field2, ...),
27+
* returns the first non-null, non-missing value in the parameter list.
28+
*
29+
* <p>Stock {@link SqlStdOperatorTable#COALESCE} requires operands that already share a common type.
30+
* So naively rewriting the call to stock {@code COALESCE} with the original operands would reject
31+
* mixed-type inputs like {@code coalesce(VARCHAR, INTEGER, "fallback")} at plan validation.
32+
*
33+
* <p>This adapter moves the coercion step to plan time by:
34+
* <ol>
35+
* <li>Computing {@code leastRestrictive} across all operand types</li>
36+
* <li>Falling back to nullable VARCHAR when the unification returns {@code null}</li>
37+
* <li>Inserting a plain CAST on each operand whose type differs from the common type.
38+
* Calcite's CAST covers numeric / boolean / character / date / time / timestamp
39+
* / decimal conversions</li>
40+
* <li>Rebuilding the call with {@link SqlStdOperatorTable#COALESCE} so isthmus serialises
41+
* it through the default Substrait catalog</li>
42+
* </ol>
43+
*
44+
* @opensearch.internal
45+
*/
46+
class CoalesceAdapter implements ScalarFunctionAdapter {
47+
48+
@Override
49+
public RexNode adapt(RexCall original, List<FieldStorageInfo> fieldStorage, RelOptCluster cluster) {
50+
RexBuilder rexBuilder = cluster.getRexBuilder();
51+
RelDataTypeFactory typeFactory = cluster.getTypeFactory();
52+
List<RexNode> operands = original.getOperands();
53+
if (operands.isEmpty()) {
54+
return original;
55+
}
56+
57+
// leastRestrictive([T]) is T, so single-operand coalesce degenerates to that operand's
58+
// type and no CAST is needed. leastRestrictive across mixed numeric types returns the
59+
// widest; across string + numeric returns VARCHAR; across incompatible type families
60+
// (e.g. DATE + BIGINT) returns null, which we backstop with VARCHAR
61+
List<RelDataType> operandTypes = operands.stream().map(RexNode::getType).toList();
62+
RelDataType commonType = typeFactory.leastRestrictive(operandTypes);
63+
if (commonType == null) {
64+
commonType = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.VARCHAR), true);
65+
}
66+
67+
List<RexNode> coerced = new ArrayList<>(operands.size());
68+
boolean anyCast = false;
69+
for (RexNode operand : operands) {
70+
if (operand.getType().equals(commonType)) {
71+
coerced.add(operand);
72+
} else {
73+
coerced.add(rexBuilder.makeCast(commonType, operand, true, false));
74+
anyCast = true;
75+
}
76+
}
77+
List<RexNode> finalOperands = anyCast ? coerced : operands;
78+
return rexBuilder.makeCall(original.getType(), SqlStdOperatorTable.COALESCE, finalOperands);
79+
}
80+
}

‎sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,9 @@ public class DataFusionAnalyticsBackendPlugin implements AnalyticsSearchBackendP
123123
// rejects the operator with "No backend supports scalar function [CASE] among [datafusion]"
124124
// before substrait emission.
125125
ScalarFunction.CASE,
126+
ScalarFunction.IS_NULL,
127+
ScalarFunction.IS_NOT_NULL,
128+
ScalarFunction.NULLIF,
126129
// ABS / SUBSTRING — PPL sort-pushdown moves these into the project tree; DataFusion has
127130
// both natively and isthmus's default catalog binds them, so no adapter needed.
128131
ScalarFunction.ABS,
@@ -477,6 +480,7 @@ public Map<ScalarFunction, ScalarFunctionAdapter> scalarFunctionAdapters() {
477480
Map.entry(ScalarFunction.MVZIP, new MvzipAdapter()),
478481
Map.entry(ScalarFunction.MVAPPEND, new MvappendAdapter()),
479482
Map.entry(ScalarFunction.BINARY, new BinaryFunctionAdapter()),
483+
Map.entry(ScalarFunction.COALESCE, new CoalesceAdapter()),
480484
Map.entry(ScalarFunction.CONCAT, new ConcatFunctionAdapter()),
481485
Map.entry(ScalarFunction.CONVERT_TZ, new ConvertTzAdapter()),
482486
Map.entry(ScalarFunction.COSH, new HyperbolicOperatorAdapter(SqlLibraryOperators.COSH)),

‎sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/CoalesceAdapterTests.java‎

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.be.datafusion;
10+
11+
import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
12+
import org.apache.calcite.plan.RelOptCluster;
13+
import org.apache.calcite.plan.hep.HepPlanner;
14+
import org.apache.calcite.plan.hep.HepProgramBuilder;
15+
import org.apache.calcite.rel.type.RelDataType;
16+
import org.apache.calcite.rel.type.RelDataTypeFactory;
17+
import org.apache.calcite.rex.RexBuilder;
18+
import org.apache.calcite.rex.RexCall;
19+
import org.apache.calcite.rex.RexNode;
20+
import org.apache.calcite.sql.SqlFunction;
21+
import org.apache.calcite.sql.SqlFunctionCategory;
22+
import org.apache.calcite.sql.SqlKind;
23+
import org.apache.calcite.sql.SqlOperator;
24+
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
25+
import org.apache.calcite.sql.type.OperandTypes;
26+
import org.apache.calcite.sql.type.ReturnTypes;
27+
import org.apache.calcite.sql.type.SqlTypeName;
28+
import org.opensearch.test.OpenSearchTestCase;
29+
30+
import java.util.ArrayList;
31+
import java.util.List;
32+
33+
/**
34+
* Unit tests for {@link CoalesceAdapter}.
35+
*/
36+
public class CoalesceAdapterTests extends OpenSearchTestCase {
37+
38+
private RelDataTypeFactory typeFactory;
39+
private RexBuilder rexBuilder;
40+
private RelOptCluster cluster;
41+
42+
private RelDataType nullableVarchar;
43+
private RelDataType notNullVarchar;
44+
private RelDataType nullableBigint;
45+
46+
private final CoalesceAdapter adapter = new CoalesceAdapter();
47+
private SqlOperator coalesce;
48+
49+
@Override
50+
public void setUp() throws Exception {
51+
super.setUp();
52+
typeFactory = new JavaTypeFactoryImpl();
53+
rexBuilder = new RexBuilder(typeFactory);
54+
HepPlanner planner = new HepPlanner(new HepProgramBuilder().build());
55+
cluster = RelOptCluster.create(planner, rexBuilder);
56+
57+
nullableVarchar = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.VARCHAR), true);
58+
notNullVarchar = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.VARCHAR), false);
59+
nullableBigint = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true);
60+
coalesce = new SqlFunction(
61+
"COALESCE",
62+
SqlKind.OTHER_FUNCTION,
63+
ReturnTypes.explicit(nullableVarchar),
64+
null,
65+
OperandTypes.VARIADIC,
66+
SqlFunctionCategory.SYSTEM
67+
);
68+
}
69+
70+
/** Builds an COALESCE call over the provided operand types. */
71+
private RexCall buildCoalesce(RelDataType... operandTypes) {
72+
List<RexNode> operands = new ArrayList<>(operandTypes.length);
73+
for (int i = 0; i < operandTypes.length; i++) {
74+
operands.add(rexBuilder.makeInputRef(operandTypes[i], i));
75+
}
76+
RelDataType returnType = typeFactory.leastRestrictive(List.of(operandTypes));
77+
if (returnType == null) {
78+
returnType = nullableVarchar;
79+
}
80+
return (RexCall) rexBuilder.makeCall(returnType, coalesce, operands);
81+
}
82+
83+
// ── operator identity rewrite ──────────────────────────────────────────
84+
85+
public void testAdaptRewritesOperatorToStockCoalesce() {
86+
RexCall original = buildCoalesce(nullableVarchar, nullableVarchar);
87+
RexNode adapted = adapter.adapt(original, List.of(), cluster);
88+
89+
assertTrue("expected RexCall, got " + adapted.getClass().getSimpleName(), adapted instanceof RexCall);
90+
RexCall rewritten = (RexCall) adapted;
91+
assertSame(
92+
"operator must be rewritten to stock SqlStdOperatorTable.COALESCE — isthmus' default "
93+
+ "catalog binds this constant to substrait's coalesce, the UDF binding is opaque",
94+
SqlStdOperatorTable.COALESCE,
95+
rewritten.getOperator()
96+
);
97+
}
98+
99+
// ── type preservation ──────────────────────────────────────────────────
100+
101+
public void testAdaptPreservesOriginalReturnType() {
102+
RexCall original = buildCoalesce(nullableVarchar, nullableVarchar);
103+
RexNode adapted = adapter.adapt(original, List.of(), cluster);
104+
105+
assertEquals(
106+
"rewritten COALESCE must declare the same return type as the original ENHANCED_COALESCE",
107+
original.getType(),
108+
adapted.getType()
109+
);
110+
}
111+
112+
// ── uniform-typed happy path ───────────────────────────────────────────
113+
114+
public void testAdaptOnUniformTypeOperandsSkipsCast() {
115+
RexCall original = buildCoalesce(nullableVarchar, nullableVarchar, nullableVarchar);
116+
RexCall rewritten = (RexCall) adapter.adapt(original, List.of(), cluster);
117+
118+
// No CAST inserted: operand list must reference-equal the original operands.
119+
assertEquals(3, rewritten.getOperands().size());
120+
for (int i = 0; i < 3; i++) {
121+
assertSame(
122+
"uniform-typed operand " + i + " must pass through without CAST",
123+
original.getOperands().get(i),
124+
rewritten.getOperands().get(i)
125+
);
126+
}
127+
}
128+
129+
// ── mixed-type coercion ────────────────────────────────────────────────
130+
131+
public void testAdaptOnMixedTypeOperandsInsertsCastToLeastRestrictive() {
132+
// VARCHAR + BIGINT + VARCHAR: leastRestrictive → VARCHAR (string wins over numeric in
133+
// Calcite's common-type resolution). The BIGINT operand must get a CAST to VARCHAR
134+
RexCall original = buildCoalesce(nullableVarchar, nullableBigint, nullableVarchar);
135+
RexCall rewritten = (RexCall) adapter.adapt(original, List.of(), cluster);
136+
137+
assertEquals(3, rewritten.getOperands().size());
138+
assertSame(
139+
"VARCHAR operand [0] already matches leastRestrictive — no CAST needed",
140+
original.getOperands().get(0),
141+
rewritten.getOperands().get(0)
142+
);
143+
144+
RexNode coercedMiddle = rewritten.getOperands().get(1);
145+
assertTrue("BIGINT operand must be wrapped in CAST", coercedMiddle instanceof RexCall);
146+
RexCall castCall = (RexCall) coercedMiddle;
147+
assertEquals("wrapper must be CAST", SqlKind.CAST, castCall.getKind());
148+
assertEquals(
149+
"CAST target must be VARCHAR family (leastRestrictive resolution of VARCHAR + BIGINT)",
150+
SqlTypeName.VARCHAR,
151+
castCall.getType().getSqlTypeName()
152+
);
153+
assertSame("CAST must wrap the original BIGINT operand", original.getOperands().get(1), castCall.getOperands().get(0));
154+
155+
assertSame(
156+
"VARCHAR operand [2] already matches leastRestrictive — no CAST needed",
157+
original.getOperands().get(2),
158+
rewritten.getOperands().get(2)
159+
);
160+
}
161+
162+
// ── nullability unification in leastRestrictive ────────────────────────
163+
164+
public void testAdaptOnNullablePlusNotNullVarcharUnifiesToNullable() {
165+
// leastRestrictive(nullable VARCHAR, NOT NULL VARCHAR) = nullable VARCHAR (nullability is
166+
// widened to accept both). The NOT NULL operand gets CAST to nullable so both operands
167+
// share a type.
168+
RexCall original = buildCoalesce(nullableVarchar, notNullVarchar);
169+
RexCall rewritten = (RexCall) adapter.adapt(original, List.of(), cluster);
170+
171+
assertEquals(2, rewritten.getOperands().size());
172+
assertSame(
173+
"already-nullable operand [0] must pass through untouched",
174+
original.getOperands().get(0),
175+
rewritten.getOperands().get(0)
176+
);
177+
178+
RexNode coerced = rewritten.getOperands().get(1);
179+
assertTrue("NOT NULL VARCHAR operand must be wrapped in CAST to nullable-VARCHAR", coerced instanceof RexCall);
180+
assertEquals(SqlKind.CAST, coerced.getKind());
181+
assertTrue("CAST target must be nullable (leastRestrictive of nullable+not-null is nullable)", coerced.getType().isNullable());
182+
}
183+
184+
// ── empty-operand pass-through ─────────────────────────────────────────
185+
186+
public void testAdaptOnEmptyOperandListPassesThrough() {
187+
RexCall original = (RexCall) rexBuilder.makeCall(nullableVarchar, coalesce, List.of());
188+
RexNode adapted = adapter.adapt(original, List.of(), cluster);
189+
190+
assertSame("empty-operand call must pass through unmodified", original, adapted);
191+
}
192+
}

0 commit comments

Comments
 (0)