Skip to content

Commit c99898a

Browse files
Fix binning udf resolution / Add type coercion support for binning UDFs (#4742) (#4780)
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
1 parent 22c7f1f commit c99898a

File tree

22 files changed

+500
-288
lines changed

22 files changed

+500
-288
lines changed

core/src/main/java/org/opensearch/sql/calcite/utils/OpenSearchTypeFactory.java

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -357,12 +357,12 @@ public static boolean isUserDefinedType(RelDataType type) {
357357
}
358358

359359
/**
360-
* Checks if the RelDataType represents a numeric field. Supports both standard SQL numeric types
361-
* (INTEGER, BIGINT, SMALLINT, TINYINT, FLOAT, DOUBLE, DECIMAL, REAL) and OpenSearch UDT numeric
362-
* types.
360+
* Checks if the RelDataType represents a numeric type. Supports standard SQL numeric types
361+
* (INTEGER, BIGINT, SMALLINT, TINYINT, FLOAT, DOUBLE, DECIMAL, REAL), OpenSearch UDT numeric
362+
* types, and string types (VARCHAR, CHAR).
363363
*
364364
* @param fieldType the RelDataType to check
365-
* @return true if the type is numeric, false otherwise
365+
* @return true if the type is numeric or string, false otherwise
366366
*/
367367
public static boolean isNumericType(RelDataType fieldType) {
368368
// Check standard SQL numeric types
@@ -378,6 +378,11 @@ public static boolean isNumericType(RelDataType fieldType) {
378378
return true;
379379
}
380380

381+
// Check string types (VARCHAR, CHAR)
382+
if (sqlType == SqlTypeName.VARCHAR || sqlType == SqlTypeName.CHAR) {
383+
return true;
384+
}
385+
381386
// Check for OpenSearch UDT numeric types
382387
if (isUserDefinedType(fieldType)) {
383388
AbstractExprRelDataType<?> exprType = (AbstractExprRelDataType<?>) fieldType;

core/src/main/java/org/opensearch/sql/calcite/utils/PPLOperandTypes.java

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,65 @@ private PPLOperandTypes() {}
136136
SqlTypeFamily.NUMERIC,
137137
SqlTypeFamily.NUMERIC,
138138
SqlTypeFamily.NUMERIC));
139+
140+
public static final UDFOperandMetadata WIDTH_BUCKET_OPERAND =
141+
UDFOperandMetadata.wrap(
142+
(CompositeOperandTypeChecker)
143+
// 1. Numeric fields: bin age span=10
144+
OperandTypes.family(
145+
SqlTypeFamily.NUMERIC,
146+
SqlTypeFamily.INTEGER,
147+
SqlTypeFamily.NUMERIC,
148+
SqlTypeFamily.NUMERIC)
149+
// 2. Timestamp fields with OpenSearch type system
150+
// Used in: Production + Integration tests (CalciteBinCommandIT)
151+
.or(
152+
OperandTypes.family(
153+
SqlTypeFamily.TIMESTAMP,
154+
SqlTypeFamily.INTEGER,
155+
SqlTypeFamily.CHARACTER, // TIMESTAMP - TIMESTAMP = INTERVAL (as STRING)
156+
SqlTypeFamily.TIMESTAMP))
157+
// 3. Timestamp fields with Calcite SCOTT schema
158+
// Used in: Unit tests (CalcitePPLBinTest)
159+
.or(
160+
OperandTypes.family(
161+
SqlTypeFamily.TIMESTAMP,
162+
SqlTypeFamily.INTEGER,
163+
SqlTypeFamily.TIMESTAMP, // TIMESTAMP - TIMESTAMP = TIMESTAMP
164+
SqlTypeFamily.TIMESTAMP))
165+
// DATE field with OpenSearch type system
166+
// Used in: Production + Integration tests (CalciteBinCommandIT)
167+
.or(
168+
OperandTypes.family(
169+
SqlTypeFamily.DATE,
170+
SqlTypeFamily.INTEGER,
171+
SqlTypeFamily.CHARACTER, // DATE - DATE = INTERVAL (as STRING)
172+
SqlTypeFamily.DATE))
173+
// DATE field with Calcite SCOTT schema
174+
// Used in: Unit tests (CalcitePPLBinTest)
175+
.or(
176+
OperandTypes.family(
177+
SqlTypeFamily.DATE,
178+
SqlTypeFamily.INTEGER,
179+
SqlTypeFamily.DATE, // DATE - DATE = DATE
180+
SqlTypeFamily.DATE))
181+
// TIME field with OpenSearch type system
182+
// Used in: Production + Integration tests (CalciteBinCommandIT)
183+
.or(
184+
OperandTypes.family(
185+
SqlTypeFamily.TIME,
186+
SqlTypeFamily.INTEGER,
187+
SqlTypeFamily.CHARACTER, // TIME - TIME = INTERVAL (as STRING)
188+
SqlTypeFamily.TIME))
189+
// TIME field with Calcite SCOTT schema
190+
// Used in: Unit tests (CalcitePPLBinTest)
191+
.or(
192+
OperandTypes.family(
193+
SqlTypeFamily.TIME,
194+
SqlTypeFamily.INTEGER,
195+
SqlTypeFamily.TIME, // TIME - TIME = TIME
196+
SqlTypeFamily.TIME)));
197+
139198
public static final UDFOperandMetadata NUMERIC_NUMERIC_NUMERIC_NUMERIC_NUMERIC =
140199
UDFOperandMetadata.wrap(
141200
OperandTypes.family(

core/src/main/java/org/opensearch/sql/calcite/utils/binning/BinnableField.java

Lines changed: 5 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,7 @@
1111
import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
1212
import org.opensearch.sql.exception.SemanticCheckException;
1313

14-
/**
15-
* Represents a validated field that supports binning operations. The existence of this class
16-
* guarantees that validation has been run - the field is either numeric or time-based.
17-
*
18-
* <p>This design encodes validation in the type system, preventing downstream code from forgetting
19-
* to validate or running validation multiple times.
20-
*/
14+
/** Represents a field that supports binning operations. */
2115
@Getter
2216
public class BinnableField {
2317
private final RexNode fieldExpr;
@@ -27,13 +21,12 @@ public class BinnableField {
2721
private final boolean isNumeric;
2822

2923
/**
30-
* Creates a validated BinnableField. Throws SemanticCheckException if the field is neither
31-
* numeric nor time-based.
24+
* Creates a BinnableField. Validates that the field type is compatible with binning operations.
3225
*
3326
* @param fieldExpr The Rex expression for the field
3427
* @param fieldType The relational data type of the field
3528
* @param fieldName The name of the field (for error messages)
36-
* @throws SemanticCheckException if the field is neither numeric nor time-based
29+
* @throws SemanticCheckException if the field type is not supported for binning
3730
*/
3831
public BinnableField(RexNode fieldExpr, RelDataType fieldType, String fieldName) {
3932
this.fieldExpr = fieldExpr;
@@ -43,13 +36,10 @@ public BinnableField(RexNode fieldExpr, RelDataType fieldType, String fieldName)
4336
this.isTimeBased = OpenSearchTypeFactory.isTimeBasedType(fieldType);
4437
this.isNumeric = OpenSearchTypeFactory.isNumericType(fieldType);
4538

46-
// Validation: field must be either numeric or time-based
39+
// Reject truly unsupported types (e.g., BOOLEAN, ARRAY, MAP)
4740
if (!isNumeric && !isTimeBased) {
4841
throw new SemanticCheckException(
49-
String.format(
50-
"Cannot apply binning: field '%s' is non-numeric and not time-related, expected"
51-
+ " numeric or time-related type",
52-
fieldName));
42+
String.format("Cannot apply binning to field '%s': unsupported type", fieldName));
5343
}
5444
}
5545

core/src/main/java/org/opensearch/sql/calcite/utils/binning/handlers/CountBinHandler.java

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
package org.opensearch.sql.calcite.utils.binning.handlers;
77

88
import org.apache.calcite.rex.RexNode;
9-
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
109
import org.opensearch.sql.ast.expression.Literal;
1110
import org.opensearch.sql.ast.tree.Bin;
1211
import org.opensearch.sql.ast.tree.CountBin;
@@ -16,7 +15,8 @@
1615
import org.opensearch.sql.calcite.utils.binning.BinFieldValidator;
1716
import org.opensearch.sql.calcite.utils.binning.BinHandler;
1817
import org.opensearch.sql.calcite.utils.binning.BinnableField;
19-
import org.opensearch.sql.expression.function.PPLBuiltinOperators;
18+
import org.opensearch.sql.expression.function.BuiltinFunctionName;
19+
import org.opensearch.sql.expression.function.PPLFuncImpTable;
2020

2121
/** Handler for bins-based (count) binning operations. */
2222
public class CountBinHandler implements BinHandler {
@@ -40,7 +40,9 @@ public RexNode createExpression(
4040
// Calculate data range using window functions
4141
RexNode minValue = context.relBuilder.min(fieldExpr).over().toRex();
4242
RexNode maxValue = context.relBuilder.max(fieldExpr).over().toRex();
43-
RexNode dataRange = context.relBuilder.call(SqlStdOperatorTable.MINUS, maxValue, minValue);
43+
RexNode dataRange =
44+
PPLFuncImpTable.INSTANCE.resolve(
45+
context.rexBuilder, BuiltinFunctionName.SUBTRACT, maxValue, minValue);
4446

4547
// Convert start/end parameters
4648
RexNode startValue = convertParameter(countBin.getStart(), context);
@@ -49,8 +51,13 @@ public RexNode createExpression(
4951
// WIDTH_BUCKET(field_value, num_bins, data_range, max_value)
5052
RexNode numBins = context.relBuilder.literal(requestedBins);
5153

52-
return context.rexBuilder.makeCall(
53-
PPLBuiltinOperators.WIDTH_BUCKET, fieldExpr, numBins, dataRange, maxValue);
54+
return PPLFuncImpTable.INSTANCE.resolve(
55+
context.rexBuilder,
56+
BuiltinFunctionName.WIDTH_BUCKET,
57+
fieldExpr,
58+
numBins,
59+
dataRange,
60+
maxValue);
5461
}
5562

5663
private RexNode convertParameter(

core/src/main/java/org/opensearch/sql/calcite/utils/binning/handlers/DefaultBinHandler.java

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
import org.opensearch.sql.calcite.utils.binning.BinHandler;
1717
import org.opensearch.sql.calcite.utils.binning.BinnableField;
1818
import org.opensearch.sql.calcite.utils.binning.RangeFormatter;
19+
import org.opensearch.sql.expression.function.BuiltinFunctionName;
20+
import org.opensearch.sql.expression.function.PPLFuncImpTable;
1921

2022
/** Handler for default binning when no parameters are specified. */
2123
public class DefaultBinHandler implements BinHandler {
@@ -45,7 +47,9 @@ private RexNode createNumericDefaultBinning(RexNode fieldExpr, CalcitePlanContex
4547
// Calculate data range
4648
RexNode minValue = context.relBuilder.min(fieldExpr).over().toRex();
4749
RexNode maxValue = context.relBuilder.max(fieldExpr).over().toRex();
48-
RexNode dataRange = context.relBuilder.call(SqlStdOperatorTable.MINUS, maxValue, minValue);
50+
RexNode dataRange =
51+
PPLFuncImpTable.INSTANCE.resolve(
52+
context.rexBuilder, BuiltinFunctionName.SUBTRACT, maxValue, minValue);
4953

5054
// Calculate magnitude-based width
5155
RexNode log10Range = context.relBuilder.call(SqlStdOperatorTable.LOG10, dataRange);
@@ -60,17 +64,21 @@ private RexNode createNumericDefaultBinning(RexNode fieldExpr, CalcitePlanContex
6064
// Calculate bin value
6165
RexNode binStartValue = calculateBinValue(fieldExpr, widthInt, context);
6266
RexNode binEndValue =
63-
context.relBuilder.call(SqlStdOperatorTable.PLUS, binStartValue, widthInt);
67+
PPLFuncImpTable.INSTANCE.resolve(
68+
context.rexBuilder, BuiltinFunctionName.ADD, binStartValue, widthInt);
6469

6570
return RangeFormatter.createRangeString(binStartValue, binEndValue, context);
6671
}
6772

6873
private RexNode calculateBinValue(RexNode fieldExpr, RexNode width, CalcitePlanContext context) {
6974

70-
RexNode divided = context.relBuilder.call(SqlStdOperatorTable.DIVIDE, fieldExpr, width);
75+
RexNode divided =
76+
PPLFuncImpTable.INSTANCE.resolve(
77+
context.rexBuilder, BuiltinFunctionName.DIVIDE, fieldExpr, width);
7178

7279
RexNode floored = context.relBuilder.call(SqlStdOperatorTable.FLOOR, divided);
7380

74-
return context.relBuilder.call(SqlStdOperatorTable.MULTIPLY, floored, width);
81+
return PPLFuncImpTable.INSTANCE.resolve(
82+
context.rexBuilder, BuiltinFunctionName.MULTIPLY, floored, width);
7583
}
7684
}

core/src/main/java/org/opensearch/sql/calcite/utils/binning/handlers/LogSpanHelper.java

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
import org.opensearch.sql.calcite.utils.binning.BinConstants;
1212
import org.opensearch.sql.calcite.utils.binning.RangeFormatter;
1313
import org.opensearch.sql.calcite.utils.binning.SpanInfo;
14+
import org.opensearch.sql.expression.function.BuiltinFunctionName;
15+
import org.opensearch.sql.expression.function.PPLFuncImpTable;
1416

1517
/** Helper for creating logarithmic span expressions. */
1618
public class LogSpanHelper {
@@ -31,14 +33,19 @@ public RexNode createLogSpanExpression(
3133
RexNode adjustedField = fieldExpr;
3234
if (coefficient != 1.0) {
3335
adjustedField =
34-
context.relBuilder.call(
35-
SqlStdOperatorTable.DIVIDE, fieldExpr, context.relBuilder.literal(coefficient));
36+
PPLFuncImpTable.INSTANCE.resolve(
37+
context.rexBuilder,
38+
BuiltinFunctionName.DIVIDE,
39+
fieldExpr,
40+
context.relBuilder.literal(coefficient));
3641
}
3742

3843
// Calculate log_base(adjusted_field)
3944
RexNode lnField = context.relBuilder.call(SqlStdOperatorTable.LN, adjustedField);
4045
RexNode lnBase = context.relBuilder.literal(Math.log(base));
41-
RexNode logValue = context.relBuilder.call(SqlStdOperatorTable.DIVIDE, lnField, lnBase);
46+
RexNode logValue =
47+
PPLFuncImpTable.INSTANCE.resolve(
48+
context.rexBuilder, BuiltinFunctionName.DIVIDE, lnField, lnBase);
4249

4350
// Get bin number
4451
RexNode binNumber = context.relBuilder.call(SqlStdOperatorTable.FLOOR, logValue);
@@ -49,15 +56,20 @@ public RexNode createLogSpanExpression(
4956

5057
RexNode basePowerBin = context.relBuilder.call(SqlStdOperatorTable.POWER, baseNode, binNumber);
5158
RexNode lowerBound =
52-
context.relBuilder.call(SqlStdOperatorTable.MULTIPLY, coefficientNode, basePowerBin);
59+
PPLFuncImpTable.INSTANCE.resolve(
60+
context.rexBuilder, BuiltinFunctionName.MULTIPLY, coefficientNode, basePowerBin);
5361

5462
RexNode binPlusOne =
55-
context.relBuilder.call(
56-
SqlStdOperatorTable.PLUS, binNumber, context.relBuilder.literal(1.0));
63+
PPLFuncImpTable.INSTANCE.resolve(
64+
context.rexBuilder,
65+
BuiltinFunctionName.ADD,
66+
binNumber,
67+
context.relBuilder.literal(1.0));
5768
RexNode basePowerBinPlusOne =
5869
context.relBuilder.call(SqlStdOperatorTable.POWER, baseNode, binPlusOne);
5970
RexNode upperBound =
60-
context.relBuilder.call(SqlStdOperatorTable.MULTIPLY, coefficientNode, basePowerBinPlusOne);
71+
PPLFuncImpTable.INSTANCE.resolve(
72+
context.rexBuilder, BuiltinFunctionName.MULTIPLY, coefficientNode, basePowerBinPlusOne);
6173

6274
// Create range string
6375
RexNode rangeStr = RangeFormatter.createRangeString(lowerBound, upperBound, context);

core/src/main/java/org/opensearch/sql/calcite/utils/binning/handlers/MinSpanBinHandler.java

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99

1010
import org.apache.calcite.rex.RexLiteral;
1111
import org.apache.calcite.rex.RexNode;
12-
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
1312
import org.opensearch.sql.ast.expression.Literal;
1413
import org.opensearch.sql.ast.tree.Bin;
1514
import org.opensearch.sql.ast.tree.MinSpanBin;
@@ -18,7 +17,8 @@
1817
import org.opensearch.sql.calcite.utils.binning.BinFieldValidator;
1918
import org.opensearch.sql.calcite.utils.binning.BinHandler;
2019
import org.opensearch.sql.calcite.utils.binning.BinnableField;
21-
import org.opensearch.sql.expression.function.PPLBuiltinOperators;
20+
import org.opensearch.sql.expression.function.BuiltinFunctionName;
21+
import org.opensearch.sql.expression.function.PPLFuncImpTable;
2222

2323
/** Handler for minspan-based binning operations. */
2424
public class MinSpanBinHandler implements BinHandler {
@@ -51,7 +51,9 @@ public RexNode createExpression(
5151
// Calculate data range using window functions
5252
RexNode minValue = context.relBuilder.min(fieldExpr).over().toRex();
5353
RexNode maxValue = context.relBuilder.max(fieldExpr).over().toRex();
54-
RexNode dataRange = context.relBuilder.call(SqlStdOperatorTable.MINUS, maxValue, minValue);
54+
RexNode dataRange =
55+
PPLFuncImpTable.INSTANCE.resolve(
56+
context.rexBuilder, BuiltinFunctionName.SUBTRACT, maxValue, minValue);
5557

5658
// Convert start/end parameters
5759
RexNode startValue = convertParameter(minSpanBin.getStart(), context);
@@ -60,8 +62,13 @@ public RexNode createExpression(
6062
// MINSPAN_BUCKET(field_value, min_span, data_range, max_value)
6163
RexNode minSpanParam = context.relBuilder.literal(minspan);
6264

63-
return context.rexBuilder.makeCall(
64-
PPLBuiltinOperators.MINSPAN_BUCKET, fieldExpr, minSpanParam, dataRange, maxValue);
65+
return PPLFuncImpTable.INSTANCE.resolve(
66+
context.rexBuilder,
67+
BuiltinFunctionName.MINSPAN_BUCKET,
68+
fieldExpr,
69+
minSpanParam,
70+
dataRange,
71+
maxValue);
6572
}
6673

6774
private RexNode convertParameter(

core/src/main/java/org/opensearch/sql/calcite/utils/binning/handlers/NumericSpanHelper.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77

88
import org.apache.calcite.rex.RexNode;
99
import org.opensearch.sql.calcite.CalcitePlanContext;
10-
import org.opensearch.sql.expression.function.PPLBuiltinOperators;
10+
import org.opensearch.sql.expression.function.BuiltinFunctionName;
11+
import org.opensearch.sql.expression.function.PPLFuncImpTable;
1112

1213
/** Helper for creating numeric span expressions. */
1314
public class NumericSpanHelper {
@@ -32,6 +33,7 @@ private RexNode createExpression(
3233
RexNode fieldExpr, RexNode spanValue, CalcitePlanContext context) {
3334

3435
// SPAN_BUCKET(field_value, span_value)
35-
return context.rexBuilder.makeCall(PPLBuiltinOperators.SPAN_BUCKET, fieldExpr, spanValue);
36+
return PPLFuncImpTable.INSTANCE.resolve(
37+
context.rexBuilder, BuiltinFunctionName.SPAN_BUCKET, fieldExpr, spanValue);
3638
}
3739
}

core/src/main/java/org/opensearch/sql/calcite/utils/binning/handlers/RangeBinHandler.java

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@
1313
import org.opensearch.sql.calcite.utils.binning.BinFieldValidator;
1414
import org.opensearch.sql.calcite.utils.binning.BinHandler;
1515
import org.opensearch.sql.calcite.utils.binning.BinnableField;
16-
import org.opensearch.sql.expression.function.PPLBuiltinOperators;
16+
import org.opensearch.sql.expression.function.BuiltinFunctionName;
17+
import org.opensearch.sql.expression.function.PPLFuncImpTable;
1718

1819
/** Handler for range-based binning (start/end parameters only). */
1920
public class RangeBinHandler implements BinHandler {
@@ -43,8 +44,14 @@ public RexNode createExpression(
4344
RexNode endParam = convertParameter(rangeBin.getEnd(), context, visitor);
4445

4546
// Use RANGE_BUCKET with data bounds and user parameters
46-
return context.rexBuilder.makeCall(
47-
PPLBuiltinOperators.RANGE_BUCKET, fieldExpr, dataMin, dataMax, startParam, endParam);
47+
return PPLFuncImpTable.INSTANCE.resolve(
48+
context.rexBuilder,
49+
BuiltinFunctionName.RANGE_BUCKET,
50+
fieldExpr,
51+
dataMin,
52+
dataMax,
53+
startParam,
54+
endParam);
4855
}
4956

5057
private RexNode convertParameter(

0 commit comments

Comments
 (0)