You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "chenkovsky (via GitHub)" <gi...@apache.org> on 2025/09/19 16:37:43 UTC
[PR] feat: spark udf array shuffle [datafusion]
chenkovsky opened a new pull request, #17674:
URL: https://github.com/apache/datafusion/pull/17674
## Which issue does this PR close?
## Rationale for this change
support shuffle udf
## What changes are included in this PR?
support shuffle udf
## Are these changes tested?
UT
## Are there any user-facing changes?
No
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: spark udf array shuffle [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on PR #17674:
URL: https://github.com/apache/datafusion/pull/17674#issuecomment-3388049897
Are we planning to implement the seed argument? There wasn't any response to my comments 🙁 @chenkovsky
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: spark udf array shuffle [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on PR #17674:
URL: https://github.com/apache/datafusion/pull/17674#issuecomment-3388333532
> > Are we planning to implement the seed argument? There wasn't any response to my comments 🙁 @chenkovsky
>
> @Jefffrey , for seed, it's more complex. In spark, different partition have different seeds. (real_seed = partition_id + user_specified_seed). In datafusion, I think we should use (real_seed = partition_id + record_batch_idx_in_partition + user_specified_seed). But currently, there's no way to get partition_id or record_batch_id. If we don't put partition_id and record_batch_idx_in_partition into seed, all record batches will have same shuffle behavior. I think it's incorrect. I created another issue. #17686
So are we fine with merging this Spark function without implementing seed argument? If this was ported from sail or comet I assume they were fine with this, though I see no mention of if this was the case. So long as we track this in an issue (e.g. in the referenced issue) I guess it's fine for now.
My only concern other than that is this could make the tests flaky, as some of the tests rely on the shuffled output being not equal to sorted output.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: spark udf array shuffle [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey commented on code in PR #17674:
URL: https://github.com/apache/datafusion/pull/17674#discussion_r2364989036
##########
datafusion/spark/src/function/array/shuffle.rs:
##########
@@ -0,0 +1,191 @@
+// 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 crate::function::functions_nested_utils::make_scalar_function;
+use arrow::array::{
+ Array, ArrayRef, Capacities, FixedSizeListArray, GenericListArray, MutableArrayData,
+ OffsetSizeTrait,
+};
+use arrow::buffer::OffsetBuffer;
+use arrow::datatypes::DataType::{FixedSizeList, LargeList, List, Null};
+use arrow::datatypes::{DataType, FieldRef};
+use datafusion_common::cast::{
+ as_fixed_size_list_array, as_large_list_array, as_list_array,
+};
+use datafusion_common::{exec_err, utils::take_function_args, Result};
+use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
+use rand::rng;
+use rand::seq::SliceRandom;
+use std::any::Any;
+use std::sync::Arc;
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkShuffle {
+ signature: Signature,
+}
+
+impl Default for SparkShuffle {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkShuffle {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::any(1, Volatility::Volatile),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkShuffle {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "shuffle"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ Ok(arg_types[0].clone())
+ }
+
+ fn invoke_with_args(
+ &self,
+ args: datafusion_expr::ScalarFunctionArgs,
+ ) -> Result<ColumnarValue> {
+ make_scalar_function(array_shuffle_inner)(&args.args)
+ }
+}
+
+/// array_shuffle SQL function
+pub fn array_shuffle_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
+ let [input_array] = take_function_args("shuffle", arg)?;
+ match &input_array.data_type() {
+ List(field) => {
+ let array = as_list_array(input_array)?;
+ general_array_shuffle::<i32>(array, field)
+ }
+ LargeList(field) => {
+ let array = as_large_list_array(input_array)?;
+ general_array_shuffle::<i64>(array, field)
+ }
+ FixedSizeList(field, _) => {
+ let array = as_fixed_size_list_array(input_array)?;
+ fixed_size_array_shuffle(array, field)
+ }
+ Null => Ok(Arc::clone(input_array)),
+ array_type => exec_err!("shuffle does not support type '{array_type}'."),
+ }
+}
+
+fn general_array_shuffle<O: OffsetSizeTrait + TryFrom<i64>>(
Review Comment:
```suggestion
fn general_array_shuffle<O: OffsetSizeTrait>(
```
##########
datafusion/sqllogictest/test_files/spark/array/shuffle.slt:
##########
Review Comment:
I think if we implement with a seed argument we can have deterministic tests for shuffle, without running it through sort or relying on the shuffled permutation not being equal to the sorted version
##########
datafusion/spark/src/function/array/shuffle.rs:
##########
@@ -0,0 +1,191 @@
+// 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 crate::function::functions_nested_utils::make_scalar_function;
+use arrow::array::{
+ Array, ArrayRef, Capacities, FixedSizeListArray, GenericListArray, MutableArrayData,
+ OffsetSizeTrait,
+};
+use arrow::buffer::OffsetBuffer;
+use arrow::datatypes::DataType::{FixedSizeList, LargeList, List, Null};
+use arrow::datatypes::{DataType, FieldRef};
+use datafusion_common::cast::{
+ as_fixed_size_list_array, as_large_list_array, as_list_array,
+};
+use datafusion_common::{exec_err, utils::take_function_args, Result};
+use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
+use rand::rng;
+use rand::seq::SliceRandom;
+use std::any::Any;
+use std::sync::Arc;
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkShuffle {
+ signature: Signature,
+}
+
+impl Default for SparkShuffle {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkShuffle {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::any(1, Volatility::Volatile),
Review Comment:
```suggestion
signature: Signature::arrays(1, None, Volatility::Volatile),
```
Example:
https://github.com/apache/datafusion/blob/44cd972546eaf8e946fa9d4576abade7daf7efab/datafusion/functions-nested/src/empty.rs#L72-L79
(using `arrays()` instead of `array()` to avoid the coercion from `FixedSizeList` to `List`)
Although, looking at the [Spark doc](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.shuffle.html) it says it accepts an optional seed argument; do we need to include that here?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: spark udf array shuffle [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #17674:
URL: https://github.com/apache/datafusion/pull/17674#issuecomment-3392344939
I agree -- let's merge this in without the seed argument for now, and if/when we need support for seeding, we can figure out how to plumb the seen through
Thanks again @Jefffrey and @chenkovsky
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: spark udf array shuffle [datafusion]
Posted by "Jefffrey (via GitHub)" <gi...@apache.org>.
Jefffrey merged PR #17674:
URL: https://github.com/apache/datafusion/pull/17674
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: spark udf array shuffle [datafusion]
Posted by "chenkovsky (via GitHub)" <gi...@apache.org>.
chenkovsky commented on PR #17674:
URL: https://github.com/apache/datafusion/pull/17674#issuecomment-3388202127
> Are we planning to implement the seed argument? There wasn't any response to my comments 🙁 @chenkovsky
@Jefffrey , for seed, it's more complex. In spark, different partition have different seeds. (real_seed = partition_id + user_specified_seed). In datafusion, I think we should use (real_seed = partition_id + record_batch_idx_in_partition + user_specified_seed). But currently, there's no way to get partition_id or record_batch_id.
If we don't put partition_id and record_batch_idx_in_partition into seed, all record batches will have same shuffle behavior. I think it's incorrect. I created another issue. #17686
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org