Skip to content

Commit 56436fe

Browse files
grorge123grorge123
committed
feat: support NullType output types in codegen dispatch
Untyped constructors such as map(), map('a', NULL) and array() leave NullType children in their output type, and the codegen dispatch gate rejected any output type containing NullType, so the whole operator fell back to Spark. Make the type gate asymmetric: canHandle now accepts NullType (top-level or nested in array/struct/map) for the output type while still rejecting it for BoundReference inputs, since CometScalaUDFCodegen.specFor cannot build an ArrowColumnSpec for a NullVector. The output emitter maps NullType to NullVector and writes it with setNull only. Also update the Scala/Java UDF guide: NullType arguments remain unsupported, NullType return types are now supported; CalendarIntervalType is removed from the unsupported list since it has been supported since Closes #5525 Assisted-by: Claude Code (claude-fable-5)
1 parent 569eaa5 commit 56436fe

13 files changed

Lines changed: 209 additions & 6 deletions

File tree

‎docs/source/user-guide/latest/scala_java_udfs.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ This feature is enabled by default. Set `spark.comet.exec.scalaUDF.codegen.enabl
4545
- Table UDFs and generators.
4646
- Python `@udf` and Pandas `@pandas_udf`.
4747
- Hive `GenericUDF` and `SimpleUDF`.
48-
- `NullType` and `UserDefinedType` arguments and return types. UDT-typed columns fall back to Spark; to keep execution in the Comet pipeline, store and read the underlying representation directly (e.g. write MLlib `Vector` outputs as `Struct<type: Byte, size: Int, indices: Array<Int>, values: Array<Double>>` rather than `VectorUDT`).
48+
- `UserDefinedType` arguments and return types, and `NullType` arguments. UDT-typed columns fall back to Spark; to keep execution in the Comet pipeline, store and read the underlying representation directly (e.g. write MLlib `Vector` outputs as `Struct<type: Byte, size: Int, indices: Array<Int>, values: Array<Double>>` rather than `VectorUDT`). A `NullType` *return* type is supported: Comet writes an all-null Arrow vector for it.
4949
- Trees whose total nested-field count (output plus all input columns the UDF tree references) exceeds `spark.sql.codegen.maxFields` (default 100). Comet refuses these at plan time and the operator falls back to Spark.
5050

5151
When a UDF is rejected, the reason surfaces through Comet's standard fallback diagnostics; the query still runs on Spark.

‎spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -102,22 +102,31 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
102102
* operator against the Catalyst `DataType`, which keeps its collation, and does not depend on
103103
* what this predicate admits. Rejecting collated strings here would force a full Spark fallback
104104
* for a route that is already correct.
105+
*
106+
* `NullType` is output-only: [[CometBatchKernelCodegenOutput]] can write an all-null Arrow
107+
* `NullVector`, but `CometScalaUDFCodegen.specFor` cannot build an [[ArrowColumnSpec]] for one,
108+
* so a `NullType` input (nested or not) has to keep falling back to Spark.
105109
*/
106-
def isSupportedDataType(dt: DataType): Boolean = dt match {
110+
def isSupportedDataType(dt: DataType): Boolean = isSupportedDataType(dt, allowNullType = false)
111+
112+
private def isSupportedDataType(dt: DataType, allowNullType: Boolean): Boolean = dt match {
113+
case NullType => allowNullType
107114
case BooleanType | ByteType | ShortType | IntegerType | LongType => true
108115
case FloatType | DoubleType => true
109116
case _: DecimalType => true
110117
case _: StringType | _: BinaryType => true
111118
case DateType | TimestampType | TimestampNTZType => true
112119
case dt if isTimeType(dt) => true
113120
case _: YearMonthIntervalType | _: DayTimeIntervalType | CalendarIntervalType => true
114-
case ArrayType(inner, _) => isSupportedDataType(inner)
121+
case ArrayType(inner, _) => isSupportedDataType(inner, allowNullType)
115122
case st: StructType =>
116123
// `fieldNames` rebuilds an array on each call, so read it once.
117124
val names = st.fieldNames
118125
names.distinct.length == names.length &&
119-
st.fields.forall(f => isSupportedDataType(f.dataType))
120-
case mt: MapType => isSupportedDataType(mt.keyType) && isSupportedDataType(mt.valueType)
126+
st.fields.forall(f => isSupportedDataType(f.dataType, allowNullType))
127+
case mt: MapType =>
128+
isSupportedDataType(mt.keyType, allowNullType) &&
129+
isSupportedDataType(mt.valueType, allowNullType)
121130
case _ => false
122131
}
123132

@@ -142,7 +151,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
142151
* nested-field count on `spark.sql.codegen.maxFields`.
143152
*/
144153
def canHandle(boundExpr: Expression): Option[String] = {
145-
if (!isSupportedDataType(boundExpr.dataType)) {
154+
if (!isSupportedDataType(boundExpr.dataType, allowNullType = true)) {
146155
return Some(s"codegen dispatch: unsupported output type ${boundExpr.dataType}")
147156
}
148157
// Mirror WSCG's `spark.sql.codegen.maxFields` gate. Wide schemas blow the generated class's

‎spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,7 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
179179
case _: ArrayType => classOf[ListVector].getName
180180
case _: StructType => classOf[StructVector].getName
181181
case _: MapType => classOf[MapVector].getName
182+
case NullType => classOf[NullVector].getName
182183
case other =>
183184
throw new UnsupportedOperationException(
184185
s"CometBatchKernelCodegen.outputVectorClass: unsupported output type $other")
@@ -209,6 +210,11 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
209210
dataType: DataType,
210211
ctx: CodegenContext,
211212
nested: Boolean = false): OutputEmit = dataType match {
213+
case NullType =>
214+
// A NullType value is null by definition: nothing to read from `source`, and `NullVector`
215+
// has no data buffer. `setNull` is a no-op; the all-null semantics come from
216+
// `CometScalaUDFCodegen.evaluate`'s post-`process` `setValueCount`.
217+
OutputEmit("", s"$targetVec.setNull($idx);")
212218
case BooleanType =>
213219
val set = if (nested) "setSafe" else "set"
214220
OutputEmit("", s"$targetVec.$set($idx, $source ? 1 : 0);")
@@ -407,6 +413,7 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
407413
*/
408414
private def emitSpecializedGetterExpr(target: String, idx: String, elemType: DataType): String =
409415
elemType match {
416+
case NullType => "null"
410417
case BooleanType => s"$target.getBoolean($idx)"
411418
case ByteType => s"$target.getByte($idx)"
412419
case ShortType => s"$target.getShort($idx)"

‎spark/src/test/resources/sql-tests/expressions/array/array_sort_comparator.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,3 +42,7 @@ FROM test_array_sort
4242
-- all literals
4343
query
4444
SELECT array_sort(array(3, 1, 2), (l, r) -> CASE WHEN l < r THEN 1 WHEN l > r THEN -1 ELSE 0 END)
45+
46+
-- an untyped empty array leaves the element type as NullType
47+
query
48+
SELECT array_sort(array(), (l, r) -> 0)

‎spark/src/test/resources/sql-tests/expressions/array/transform.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,3 +80,7 @@ SELECT a FROM test_transform WHERE transform(a, x -> array(x)) = nested
8080
-- all literals (constant folding is disabled by the test harness)
8181
query
8282
SELECT transform(array(1, 2, 3), x -> x * x)
83+
84+
-- an untyped empty array leaves the element type as NullType
85+
query
86+
SELECT transform(array(), x -> x)

‎spark/src/test/resources/sql-tests/expressions/array/zip_with.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,3 +43,7 @@ SELECT zip_with(a, b, (x, y) -> struct(x AS l, y AS r)) FROM test_zip_with
4343
-- all literals
4444
query
4545
SELECT zip_with(array(1, 2), array(3, 4), (x, y) -> x * y)
46+
47+
-- untyped empty arrays leave the element type as NullType
48+
query
49+
SELECT zip_with(array(), array(), (x, y) -> x)

‎spark/src/test/resources/sql-tests/expressions/map/create_map.sql‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,3 +39,37 @@ SELECT map(1, named_struct('x', 10, 'x', 20))
3939
-- Distinct field names are unaffected.
4040
query
4141
SELECT map(1, named_struct('x', 10, 'y', 20))
42+
43+
-- ===== untyped constructors leave NullType children =====
44+
-- `map()` is MapType(NullType, NullType) and `map(k, NULL)` is MapType(_, NullType).
45+
46+
query
47+
SELECT map()
48+
49+
query
50+
SELECT map('a', NULL)
51+
52+
query
53+
SELECT map(k, NULL) FROM test_create_map
54+
55+
query
56+
SELECT id, map() FROM (SELECT explode(sequence(1, 3)) AS id)
57+
58+
query
59+
SELECT size(map()), size(map('a', NULL))
60+
61+
query
62+
SELECT map_keys(map()), map_values(map('a', NULL))
63+
64+
query
65+
SELECT array(map()), struct(map(), map('a', NULL))
66+
67+
query
68+
SELECT map('a', array(NULL)), map('a', map())
69+
70+
query
71+
SELECT map_from_arrays(array(), array())
72+
73+
-- carry the NullType children through a sort so the vector survives copy/spill paths
74+
query
75+
SELECT k, map(), map('a', NULL) FROM test_create_map ORDER BY k

‎spark/src/test/resources/sql-tests/expressions/map/map_concat.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,3 +49,7 @@ SELECT map_concat(map(1, 'x', 2, 'y'), map(3, 'z'))
4949
-- a NULL literal map makes the whole result NULL
5050
query
5151
SELECT map_concat(map('a', 1), CAST(NULL AS map<string, int>))
52+
53+
-- untyped operands leave the result as MapType(NullType, NullType)
54+
query
55+
SELECT map_concat(map(), map())

‎spark/src/test/resources/sql-tests/expressions/map/map_filter.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,3 +42,7 @@ SELECT map_filter(m, (k, v) -> v > threshold) FROM test_map_filter
4242
-- all literals
4343
query
4444
SELECT map_filter(map('a', 1, 'b', 2), (k, v) -> v > 1)
45+
46+
-- an untyped value leaves the value type as NullType
47+
query
48+
SELECT map_filter(map('a', NULL), (k, v) -> true)

‎spark/src/test/resources/sql-tests/expressions/map/map_zip_with.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,3 +39,7 @@ SELECT map_zip_with(m, n, (k, v1, v2) -> struct(v1 AS left, v2 AS right)) FROM t
3939
-- all literals
4040
query
4141
SELECT map_zip_with(map('a', 1), map('a', 2, 'b', 3), (k, v1, v2) -> coalesce(v1, 0) + coalesce(v2, 0))
42+
43+
-- untyped operands leave both key and value types as NullType
44+
query
45+
SELECT map_zip_with(map(), map(), (k, v1, v2) -> v1)

0 commit comments

Comments
 (0)