andygrove commented on code in PR #26266: URL: https://github.com/apache/datafusion/pull/26266#discussion_r4327326543
########## datafusion/spark/src/function/generator/stack.rs: ########## @@ -0,0 +1,517 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, ListArray, NullBufferBuilder, PrimitiveArray, StructArray, + UInt32Array, downcast_primitive, new_empty_array, new_null_array, +}; +use arrow::buffer::OffsetBuffer; +use arrow::compute::{concat, take}; +use arrow::datatypes::{ArrowPrimitiveType, DataType, Field, FieldRef, Fields}; +use datafusion_common::types::NativeType; +use datafusion_common::{ + Result, ScalarValue, internal_err, plan_err, resources_datafusion_err, +}; +use datafusion_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, +}; + +/// The most rows a call generates per batch. Spark streams the rows, but they are +/// materialized here, so a large `n` such as `stack(2000000000, 1)` fails instead of +/// exhausting memory. +const MAX_GENERATED_ROWS: usize = 1 << 24; + +/// Spark-compatible `stack` expression +/// +/// <https://spark.apache.org/docs/latest/api/sql/index.html#stack> +/// +/// `stack(n, expr1, ..., exprk)` separates `expr1`, ..., `exprk` into `n` rows of +/// `ceil(k / n)` columns named `col0`, `col1`, ... The values fill the rows in +/// order, and the last rows are padded with NULLs. +/// +/// Note: In Spark, `stack` is a Generator that produces rows. In DataFusion, a +/// ScalarUDF can only return one value per row, so the `n` rows are returned as a +/// `List<Struct<col0, col1, ...>>`, which is the data type of Spark's Generator. As +/// in Spark, the list is not nullable while its elements and fields are. The +/// caller (e.g. Comet, or `unnest` in SQL) is expected to unnest the list. +/// +/// As in Spark, `n` is an `INT` in `[1, 2147483647]`, and the values of a column +/// must have the same type, which is the type of its first non-NULL value. NULLs, +/// Arrow types of the same Spark type (such as `Utf8` and `Utf8View`), and integer +/// literals that fit an `INT` column are cast to it, wherever they are in the +/// column. Unlike Spark, which accepts any foldable `n`, `n` must be an integer +/// literal, because the return type is planned before expressions are folded. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkStack { + signature: Signature, +} + +impl Default for SparkStack { + fn default() -> Self { + Self::new() + } +} + +impl SparkStack { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkStack { + fn name(&self) -> &str { + "stack" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + internal_err!("return_field_from_args should be used instead") + } + + /// Validates the arguments in the order Spark's `Stack.checkInputDataTypes` does. + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> { + let (n, values) = match args.arg_fields { + [n, values @ ..] if !values.is_empty() => (n, values), + fields => { + return plan_err!( + "The `stack` requires > 1 parameters but the actual number is {}.", + fields.len() + ); + } + }; + if n.data_type() != &DataType::Int32 { + return plan_err!( + "The first parameter requires the \"INT\" type, however \"{}\" has the type \"{}\".", + n.name(), + n.data_type() + ); + } + let Some(n_value) = args.scalar_arguments.first().copied().flatten() else { + return plan_err!( + "The first parameter of `stack` must be an integer literal, however \"{}\" is not one.", + n.name() + ); + }; + let num_fields = values.len().div_ceil(parse_n(n_value)?); + let scalar_arguments = args.scalar_arguments.get(1..).unwrap_or_default(); + let fields = column_types(values, scalar_arguments, num_fields)? + .into_iter() + .enumerate() + .map(|(i, data_type)| Field::new(format!("col{i}"), data_type, true)) + .collect::<Fields>(); + Ok(Arc::new(Field::new( + self.name(), + DataType::new_list(DataType::Struct(fields), true), + false, + ))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + let ScalarFunctionArgs { + args, + number_rows, + return_field, + .. + } = args; + let DataType::List(item) = return_field.data_type() else { + return internal_err!("stack: unexpected return type {return_field}"); + }; + let [ColumnarValue::Scalar(n), values @ ..] = args.as_slice() else { + return internal_err!("stack: `n` must be a scalar"); + }; + let n = parse_n(n)?; + if values.iter().all(|v| matches!(v, ColumnarValue::Scalar(_))) { Review Comment: The cap is only checked against the real batch size on the array path. When every value is a literal, this calls `stack(.., 1)` and returns one scalar list, and the caller then expands it to `number_rows` copies through `ScalarValue::to_array_of_size` with no check. With `n = 2^20` and `number_rows = 8192`, `invoke_with_args` returns `Ok`, and the expanded result would be about 8.6 billion rows against a cap of 16.8 million. Something like `SELECT stack(1048576, 1) FROM generate_series(1, 8192)` should reach it, and so does constant folding, which turns the call into a literal at plan time. Could the `number_rows * n` check run before this shortcut? That covers execution, including Comet. Folded literals would still get through, but at least that gap would be known. ########## datafusion/spark/src/function/generator/stack.rs: ########## @@ -0,0 +1,517 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, ListArray, NullBufferBuilder, PrimitiveArray, StructArray, + UInt32Array, downcast_primitive, new_empty_array, new_null_array, +}; +use arrow::buffer::OffsetBuffer; +use arrow::compute::{concat, take}; +use arrow::datatypes::{ArrowPrimitiveType, DataType, Field, FieldRef, Fields}; +use datafusion_common::types::NativeType; +use datafusion_common::{ + Result, ScalarValue, internal_err, plan_err, resources_datafusion_err, +}; +use datafusion_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, +}; + +/// The most rows a call generates per batch. Spark streams the rows, but they are +/// materialized here, so a large `n` such as `stack(2000000000, 1)` fails instead of +/// exhausting memory. +const MAX_GENERATED_ROWS: usize = 1 << 24; + +/// Spark-compatible `stack` expression +/// +/// <https://spark.apache.org/docs/latest/api/sql/index.html#stack> +/// +/// `stack(n, expr1, ..., exprk)` separates `expr1`, ..., `exprk` into `n` rows of +/// `ceil(k / n)` columns named `col0`, `col1`, ... The values fill the rows in +/// order, and the last rows are padded with NULLs. +/// +/// Note: In Spark, `stack` is a Generator that produces rows. In DataFusion, a +/// ScalarUDF can only return one value per row, so the `n` rows are returned as a +/// `List<Struct<col0, col1, ...>>`, which is the data type of Spark's Generator. As +/// in Spark, the list is not nullable while its elements and fields are. The +/// caller (e.g. Comet, or `unnest` in SQL) is expected to unnest the list. +/// +/// As in Spark, `n` is an `INT` in `[1, 2147483647]`, and the values of a column +/// must have the same type, which is the type of its first non-NULL value. NULLs, +/// Arrow types of the same Spark type (such as `Utf8` and `Utf8View`), and integer +/// literals that fit an `INT` column are cast to it, wherever they are in the +/// column. Unlike Spark, which accepts any foldable `n`, `n` must be an integer +/// literal, because the return type is planned before expressions are folded. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkStack { + signature: Signature, +} + +impl Default for SparkStack { + fn default() -> Self { + Self::new() + } +} + +impl SparkStack { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkStack { + fn name(&self) -> &str { + "stack" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + internal_err!("return_field_from_args should be used instead") + } + + /// Validates the arguments in the order Spark's `Stack.checkInputDataTypes` does. + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> { + let (n, values) = match args.arg_fields { + [n, values @ ..] if !values.is_empty() => (n, values), + fields => { + return plan_err!( + "The `stack` requires > 1 parameters but the actual number is {}.", + fields.len() + ); + } + }; + if n.data_type() != &DataType::Int32 { + return plan_err!( + "The first parameter requires the \"INT\" type, however \"{}\" has the type \"{}\".", + n.name(), + n.data_type() + ); + } + let Some(n_value) = args.scalar_arguments.first().copied().flatten() else { + return plan_err!( + "The first parameter of `stack` must be an integer literal, however \"{}\" is not one.", + n.name() + ); + }; + let num_fields = values.len().div_ceil(parse_n(n_value)?); + let scalar_arguments = args.scalar_arguments.get(1..).unwrap_or_default(); + let fields = column_types(values, scalar_arguments, num_fields)? + .into_iter() + .enumerate() + .map(|(i, data_type)| Field::new(format!("col{i}"), data_type, true)) + .collect::<Fields>(); + Ok(Arc::new(Field::new( + self.name(), + DataType::new_list(DataType::Struct(fields), true), + false, + ))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + let ScalarFunctionArgs { + args, + number_rows, + return_field, + .. + } = args; + let DataType::List(item) = return_field.data_type() else { + return internal_err!("stack: unexpected return type {return_field}"); + }; + let [ColumnarValue::Scalar(n), values @ ..] = args.as_slice() else { + return internal_err!("stack: `n` must be a scalar"); + }; + let n = parse_n(n)?; + if values.iter().all(|v| matches!(v, ColumnarValue::Scalar(_))) { + let list = stack(values, item, n, 1)?; + return Ok(ColumnarValue::Scalar(ScalarValue::List(Arc::new(list)))); + } + Ok(ColumnarValue::Array(Arc::new(stack( + values, + item, + n, + number_rows, + )?))) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> { + // `n` is an `INT` as in Spark, while integer literals are BIGINT here. The + // values are cast to their column's type at execution, since Spark does no + // implicit casts. + let mut arg_types = arg_types.to_vec(); + if let Some(n) = arg_types.first_mut().filter(|n| n.is_integer()) { + *n = DataType::Int32; + } + Ok(arg_types) + } +} + +/// Returns the number of rows `n` to generate per input row. +fn parse_n(n: &ScalarValue) -> Result<usize> { + match n { + ScalarValue::Int32(Some(v)) if *v > 0 => Ok(*v as usize), + _ => plan_err!("The `n` must be between (0, 2147483647] (current value = {n})."), + } +} + +/// Returns the type of each output column, which is the type of its first +/// non-NULL value as in Spark's `StackCoercion`. Errors on the first value, in +/// argument order, whose type differs from its column's. +fn column_types( + arg_fields: &[FieldRef], + scalar_arguments: &[Option<&ScalarValue>], + num_fields: usize, +) -> Result<Vec<DataType>> { + // Integer literals are BIGINT here but INT in Spark, so one that fits is an INT + // in a column that has an INT value, wherever it is. Being a literal never + // changes the column type otherwise, which must not change when constant + // folding turns an expression into a literal. + let mut has_int32 = vec![false; num_fields]; + for (i, arg_field) in arg_fields.iter().enumerate() { + has_int32[i % num_fields] |= arg_field.data_type() == &DataType::Int32; + } + let mut types = vec![DataType::Null; num_fields]; + for (i, arg_field) in arg_fields.iter().enumerate() { + let col = i % num_fields; + let literal = scalar_arguments.get(i).copied().flatten(); + let actual = if has_int32[col] + && arg_field.data_type().is_integer() + && literal.is_some_and(|literal| literal.cast_to(&DataType::Int32).is_ok()) + { + &DataType::Int32 + } else { + arg_field.data_type() + }; + let expected = &mut types[col]; + if expected.is_null() { + *expected = actual.clone(); + } else if !actual.is_null() && !same_spark_type(actual, expected) { + return plan_err!( + "The data type of the column ({col}) do not have the same type: \"{expected}\" ({}) <> \"{actual}\" ({}).", + col + 1, + i + 1 + ); + } + } + Ok(types) +} + +/// Whether two Arrow types are the same Spark type. Types of one logical type, such +/// as `Utf8` and `Utf8View`, match, and list elements and map entries match +/// regardless of their field names, which Spark types do not have. As in Spark's +/// `sameType`, nullability is ignored. Review Comment: This comes from my last review, and I was wrong about it. Spark's `Stack.checkInputDataTypes` compares each value's type with the column type using `!=` (in both 3.5.8 and 4.1.1), so `containsNull`, `valueContainsNull` and `StructField.nullable` all count. `StackTypeCoercion` only retypes a `NULL` literal. Ignoring nullability now lets a query plan and then fail. The column takes the first value's type, so when that's `List(non-null Int64)` and a later value is a `List(Int64)` with a NULL element, the cast at execution fails with `Non-nullable field of ListArray "item" cannot contain nulls`. Something like `SELECT stack(2, arrow_cast([1], 'List(non-null Int64)'), [2, NULL])` should show it. The existing `Nullability does not make types differ` case passes only because `[2]` has no NULL. I'd go back to matching nullability exactly, as Spark does. If you'd rather stay lenient, the column type's nested fields would need to be nullable whenever the values differ. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
