You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "davidlghellin (via GitHub)" <gi...@apache.org> on 2025/08/26 11:40:07 UTC
[PR] feat(spark): implement Spark `width_bucket` function [datafusion]
davidlghellin opened a new pull request, #17331:
URL: https://github.com/apache/datafusion/pull/17331
## Which issue does this PR close?
part of https://github.com/apache/datafusion/issues/15914
## Rationale for this change
Migrate spark functions from https://github.com/lakehq/sail/ to datafusion engine to unify codebase
## What changes are included in this PR?
implement spark udf width_bucket
https://spark.apache.org/docs/latest/api/sql/index.html#width_bucket
## Are these changes tested?
unit-tests and sqllogictests added
## Are there any user-facing changes?
now can be called in queries
--
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): implement Spark `width_bucket` function [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #17331:
URL: https://github.com/apache/datafusion/pull/17331#discussion_r2304277268
##########
datafusion/spark/src/function/math/width_bucket.rs:
##########
@@ -0,0 +1,506 @@
+// 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::any::Any;
+use std::sync::Arc;
+
+use crate::function::error_utils::unsupported_data_types_exec_err;
+use arrow::array::{
+ Array, ArrayRef, DurationMicrosecondArray, Float64Array, IntervalYearMonthArray,
+};
+use arrow::datatypes::DataType;
+use arrow::datatypes::DataType::{
+ Duration, Float64, Int32, Interval,
+};
+use arrow::datatypes::IntervalUnit::YearMonth;
+use datafusion_common::cast::{
+ as_duration_microsecond_array, as_float64_array, as_int32_array, as_interval_ym_array,
+};
+use datafusion_common::{exec_err, Result};
+use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature};
+use datafusion_functions::utils::make_scalar_function;
+
+use arrow::array::{Int32Array, Int32Builder};
+use arrow::datatypes::TimeUnit::Microsecond;
+use datafusion_expr::Volatility::Immutable;
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkWidthBucket {
+ signature: Signature,
+}
+
+impl Default for SparkWidthBucket {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkWidthBucket {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::user_defined(Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkWidthBucket {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "width_bucket"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(Int32)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ make_scalar_function(width_bucket_kern, vec![])(&args.args)
+ }
+
+ fn coerce_types(&self, types: &[DataType]) -> Result<Vec<DataType>> {
+ use DataType::*;
+
+ let (v, lo, hi, n) = (&types[0], &types[1], &types[2], &types[3]);
+
+ let is_num = |t: &DataType| {
+ matches!(
+ t,
+ Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Decimal128(_, _)
+ )
+ };
+ let is_int = |t: &DataType| matches!(t, Int8 | Int16 | Int32 | Int64);
+
+ if is_num(v) && is_num(lo) && is_num(hi) && is_int(n) {
+ return Ok(vec![Float64, Float64, Float64, Int32]);
+ }
+
+ let all_duration = matches!(v, Duration(_))
+ && matches!(lo, Duration(_))
+ && matches!(hi, Duration(_));
+ if all_duration && is_int(n) {
+ return Ok(vec![
+ Duration(Microsecond),
+ Duration(Microsecond),
+ Duration(Microsecond),
+ Int32,
+ ]);
+ }
+
+ if matches!(v, Interval(YearMonth))
+ && matches!(lo, Interval(YearMonth))
+ && matches!(hi, Interval(YearMonth))
+ && is_int(n)
+ {
+ return Ok(vec![
+ Interval(YearMonth),
+ Interval(YearMonth),
+ Interval(YearMonth),
+ Int32,
+ ]);
+ }
+
+ exec_err!("rint expects a numeric argument, got {}", types[0])
+ }
+
+ fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {
+ // FIXME copy
Review Comment:
is it needed?
--
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): implement Spark `width_bucket` function [datafusion]
Posted by "davidlghellin (via GitHub)" <gi...@apache.org>.
davidlghellin commented on code in PR #17331:
URL: https://github.com/apache/datafusion/pull/17331#discussion_r2304571969
##########
datafusion/sqllogictest/test_files/spark/math/width_bucket.slt:
##########
@@ -21,42 +21,58 @@
# For more information, please see:
# https://github.com/apache/datafusion/issues/15914
-## Original Query: SELECT width_bucket(-0.9, 5.2, 0.5, 2);
-## PySpark 3.5.5 Result: {'width_bucket(-0.9, 5.2, 0.5, 2)': 3, 'typeof(width_bucket(-0.9, 5.2, 0.5, 2))': 'bigint', 'typeof(-0.9)': 'decimal(1,1)', 'typeof(5.2)': 'decimal(2,1)', 'typeof(0.5)': 'decimal(1,1)', 'typeof(2)': 'int'}
-#query
-#SELECT width_bucket(-0.9::decimal(1,1), 5.2::decimal(2,1), 0.5::decimal(1,1), 2::int);
-
-## Original Query: SELECT width_bucket(-2.1, 1.3, 3.4, 3);
-## PySpark 3.5.5 Result: {'width_bucket(-2.1, 1.3, 3.4, 3)': 0, 'typeof(width_bucket(-2.1, 1.3, 3.4, 3))': 'bigint', 'typeof(-2.1)': 'decimal(2,1)', 'typeof(1.3)': 'decimal(2,1)', 'typeof(3.4)': 'decimal(2,1)', 'typeof(3)': 'int'}
-#query
-#SELECT width_bucket(-2.1::decimal(2,1), 1.3::decimal(2,1), 3.4::decimal(2,1), 3::int);
-
-## Original Query: SELECT width_bucket(5.3, 0.2, 10.6, 5);
-## PySpark 3.5.5 Result: {'width_bucket(5.3, 0.2, 10.6, 5)': 3, 'typeof(width_bucket(5.3, 0.2, 10.6, 5))': 'bigint', 'typeof(5.3)': 'decimal(2,1)', 'typeof(0.2)': 'decimal(1,1)', 'typeof(10.6)': 'decimal(3,1)', 'typeof(5)': 'int'}
-#query
-#SELECT width_bucket(5.3::decimal(2,1), 0.2::decimal(1,1), 10.6::decimal(3,1), 5::int);
-
-## Original Query: SELECT width_bucket(8.1, 0.0, 5.7, 4);
-## PySpark 3.5.5 Result: {'width_bucket(8.1, 0.0, 5.7, 4)': 5, 'typeof(width_bucket(8.1, 0.0, 5.7, 4))': 'bigint', 'typeof(8.1)': 'decimal(2,1)', 'typeof(0.0)': 'decimal(1,1)', 'typeof(5.7)': 'decimal(2,1)', 'typeof(4)': 'int'}
-#query
-#SELECT width_bucket(8.1::decimal(2,1), 0.0::decimal(1,1), 5.7::decimal(2,1), 4::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)": 1, "typeof(width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10))": 'bigint', "typeof(INTERVAL '0' DAY)": 'interval day', "typeof(INTERVAL '10' DAY)": 'interval day', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '0' DAY::interval day, INTERVAL '10' DAY::interval day, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)": 1, "typeof(width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10))": 'bigint', "typeof(INTERVAL '0' YEAR)": 'interval year', "typeof(INTERVAL '10' YEAR)": 'interval year', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '0' YEAR::interval year, INTERVAL '10' YEAR::interval year, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)": 2, "typeof(width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10))": 'bigint', "typeof(INTERVAL '1' DAY)": 'interval day', "typeof(INTERVAL '0' DAY)": 'interval day', "typeof(INTERVAL '10' DAY)": 'interval day', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '1' DAY::interval day, INTERVAL '0' DAY::interval day, INTERVAL '10' DAY::interval day, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)": 2, "typeof(width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10))": 'bigint', "typeof(INTERVAL '1' YEAR)": 'interval year', "typeof(INTERVAL '0' YEAR)": 'interval year', "typeof(INTERVAL '10' YEAR)": 'interval year', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '1' YEAR::interval year, INTERVAL '0' YEAR::interval year, INTERVAL '10' YEAR::interval year, 10::int);
+query I
+SELECT width_bucket(-0.9, 5.2, 0.5, 2)
+----
+3
+
+query I
+SELECT width_bucket(-2.1, 1.3, 3.4, 3)
+----
+0
+
+query I
+SELECT width_bucket(5.3, 0.2, 10.6, 5)
+----
+3
+
+query I
+SELECT width_bucket(8.1, 0.0, 5.7, 4)
+----
+5
+
+query I
+SELECT width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)
Review Comment:
I'll take a look to verify these errors
--
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): implement Spark `width_bucket` function [datafusion]
Posted by "davidlghellin (via GitHub)" <gi...@apache.org>.
davidlghellin commented on code in PR #17331:
URL: https://github.com/apache/datafusion/pull/17331#discussion_r2307652098
##########
datafusion/spark/src/function/math/width_bucket.rs:
##########
@@ -0,0 +1,506 @@
+// 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::any::Any;
+use std::sync::Arc;
+
+use crate::function::error_utils::unsupported_data_types_exec_err;
+use arrow::array::{
+ Array, ArrayRef, DurationMicrosecondArray, Float64Array, IntervalYearMonthArray,
+};
+use arrow::datatypes::DataType;
+use arrow::datatypes::DataType::{
+ Duration, Float64, Int32, Interval,
+};
+use arrow::datatypes::IntervalUnit::YearMonth;
+use datafusion_common::cast::{
+ as_duration_microsecond_array, as_float64_array, as_int32_array, as_interval_ym_array,
+};
+use datafusion_common::{exec_err, Result};
+use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature};
+use datafusion_functions::utils::make_scalar_function;
+
+use arrow::array::{Int32Array, Int32Builder};
+use arrow::datatypes::TimeUnit::Microsecond;
+use datafusion_expr::Volatility::Immutable;
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkWidthBucket {
+ signature: Signature,
+}
+
+impl Default for SparkWidthBucket {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkWidthBucket {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::user_defined(Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkWidthBucket {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "width_bucket"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(Int32)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ make_scalar_function(width_bucket_kern, vec![])(&args.args)
+ }
+
+ fn coerce_types(&self, types: &[DataType]) -> Result<Vec<DataType>> {
+ use DataType::*;
+
+ let (v, lo, hi, n) = (&types[0], &types[1], &types[2], &types[3]);
+
+ let is_num = |t: &DataType| {
+ matches!(
+ t,
+ Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Decimal128(_, _)
+ )
+ };
+ let is_int = |t: &DataType| matches!(t, Int8 | Int16 | Int32 | Int64);
+
+ if is_num(v) && is_num(lo) && is_num(hi) && is_int(n) {
+ return Ok(vec![Float64, Float64, Float64, Int32]);
+ }
+
+ let all_duration = matches!(v, Duration(_))
+ && matches!(lo, Duration(_))
+ && matches!(hi, Duration(_));
+ if all_duration && is_int(n) {
+ return Ok(vec![
+ Duration(Microsecond),
+ Duration(Microsecond),
+ Duration(Microsecond),
+ Int32,
+ ]);
+ }
+
+ if matches!(v, Interval(YearMonth))
+ && matches!(lo, Interval(YearMonth))
+ && matches!(hi, Interval(YearMonth))
+ && is_int(n)
+ {
+ return Ok(vec![
+ Interval(YearMonth),
+ Interval(YearMonth),
+ Interval(YearMonth),
+ Int32,
+ ]);
+ }
+
+ exec_err!("rint expects a numeric argument, got {}", types[0])
+ }
+
+ fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {
+ // FIXME copy
Review Comment:
Hi, I drop the FIXME
--
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): implement Spark `width_bucket` function [datafusion]
Posted by "davidlghellin (via GitHub)" <gi...@apache.org>.
davidlghellin commented on PR #17331:
URL: https://github.com/apache/datafusion/pull/17331#issuecomment-3228862362
Thanks @comphead , I'm checking this because I don't understand the errors I'm getting on my machine. I'm not sure if it's something related to my setup or if I missed something in the configuration.
--
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): implement Spark `width_bucket` function [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead merged PR #17331:
URL: https://github.com/apache/datafusion/pull/17331
--
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): implement Spark `width_bucket` function [datafusion]
Posted by "davidlghellin (via GitHub)" <gi...@apache.org>.
davidlghellin commented on code in PR #17331:
URL: https://github.com/apache/datafusion/pull/17331#discussion_r2304559249
##########
datafusion/sqllogictest/test_files/spark/math/width_bucket.slt:
##########
@@ -21,42 +21,58 @@
# For more information, please see:
# https://github.com/apache/datafusion/issues/15914
-## Original Query: SELECT width_bucket(-0.9, 5.2, 0.5, 2);
-## PySpark 3.5.5 Result: {'width_bucket(-0.9, 5.2, 0.5, 2)': 3, 'typeof(width_bucket(-0.9, 5.2, 0.5, 2))': 'bigint', 'typeof(-0.9)': 'decimal(1,1)', 'typeof(5.2)': 'decimal(2,1)', 'typeof(0.5)': 'decimal(1,1)', 'typeof(2)': 'int'}
-#query
-#SELECT width_bucket(-0.9::decimal(1,1), 5.2::decimal(2,1), 0.5::decimal(1,1), 2::int);
-
-## Original Query: SELECT width_bucket(-2.1, 1.3, 3.4, 3);
-## PySpark 3.5.5 Result: {'width_bucket(-2.1, 1.3, 3.4, 3)': 0, 'typeof(width_bucket(-2.1, 1.3, 3.4, 3))': 'bigint', 'typeof(-2.1)': 'decimal(2,1)', 'typeof(1.3)': 'decimal(2,1)', 'typeof(3.4)': 'decimal(2,1)', 'typeof(3)': 'int'}
-#query
-#SELECT width_bucket(-2.1::decimal(2,1), 1.3::decimal(2,1), 3.4::decimal(2,1), 3::int);
-
-## Original Query: SELECT width_bucket(5.3, 0.2, 10.6, 5);
-## PySpark 3.5.5 Result: {'width_bucket(5.3, 0.2, 10.6, 5)': 3, 'typeof(width_bucket(5.3, 0.2, 10.6, 5))': 'bigint', 'typeof(5.3)': 'decimal(2,1)', 'typeof(0.2)': 'decimal(1,1)', 'typeof(10.6)': 'decimal(3,1)', 'typeof(5)': 'int'}
-#query
-#SELECT width_bucket(5.3::decimal(2,1), 0.2::decimal(1,1), 10.6::decimal(3,1), 5::int);
-
-## Original Query: SELECT width_bucket(8.1, 0.0, 5.7, 4);
-## PySpark 3.5.5 Result: {'width_bucket(8.1, 0.0, 5.7, 4)': 5, 'typeof(width_bucket(8.1, 0.0, 5.7, 4))': 'bigint', 'typeof(8.1)': 'decimal(2,1)', 'typeof(0.0)': 'decimal(1,1)', 'typeof(5.7)': 'decimal(2,1)', 'typeof(4)': 'int'}
-#query
-#SELECT width_bucket(8.1::decimal(2,1), 0.0::decimal(1,1), 5.7::decimal(2,1), 4::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)": 1, "typeof(width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10))": 'bigint', "typeof(INTERVAL '0' DAY)": 'interval day', "typeof(INTERVAL '10' DAY)": 'interval day', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '0' DAY::interval day, INTERVAL '10' DAY::interval day, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)": 1, "typeof(width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10))": 'bigint', "typeof(INTERVAL '0' YEAR)": 'interval year', "typeof(INTERVAL '10' YEAR)": 'interval year', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '0' YEAR::interval year, INTERVAL '10' YEAR::interval year, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)": 2, "typeof(width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10))": 'bigint', "typeof(INTERVAL '1' DAY)": 'interval day', "typeof(INTERVAL '0' DAY)": 'interval day', "typeof(INTERVAL '10' DAY)": 'interval day', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '1' DAY::interval day, INTERVAL '0' DAY::interval day, INTERVAL '10' DAY::interval day, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)": 2, "typeof(width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10))": 'bigint', "typeof(INTERVAL '1' YEAR)": 'interval year', "typeof(INTERVAL '0' YEAR)": 'interval year', "typeof(INTERVAL '10' YEAR)": 'interval year', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '1' YEAR::interval year, INTERVAL '0' YEAR::interval year, INTERVAL '10' YEAR::interval year, 10::int);
+query I
+SELECT width_bucket(-0.9, 5.2, 0.5, 2)
+----
+3
+
+query I
+SELECT width_bucket(-2.1, 1.3, 3.4, 3)
+----
+0
+
+query I
+SELECT width_bucket(5.3, 0.2, 10.6, 5)
+----
+3
+
+query I
+SELECT width_bucket(8.1, 0.0, 5.7, 4)
+----
+5
+
+query I
+SELECT width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)
Review Comment:
check interval
--
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): implement Spark `width_bucket` function [datafusion]
Posted by "davidlghellin (via GitHub)" <gi...@apache.org>.
davidlghellin commented on code in PR #17331:
URL: https://github.com/apache/datafusion/pull/17331#discussion_r2304394160
##########
datafusion/spark/src/function/math/width_bucket.rs:
##########
@@ -0,0 +1,506 @@
+// 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::any::Any;
+use std::sync::Arc;
+
+use crate::function::error_utils::unsupported_data_types_exec_err;
+use arrow::array::{
+ Array, ArrayRef, DurationMicrosecondArray, Float64Array, IntervalYearMonthArray,
+};
+use arrow::datatypes::DataType;
+use arrow::datatypes::DataType::{
+ Duration, Float64, Int32, Interval,
+};
+use arrow::datatypes::IntervalUnit::YearMonth;
+use datafusion_common::cast::{
+ as_duration_microsecond_array, as_float64_array, as_int32_array, as_interval_ym_array,
+};
+use datafusion_common::{exec_err, Result};
+use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature};
+use datafusion_functions::utils::make_scalar_function;
+
+use arrow::array::{Int32Array, Int32Builder};
+use arrow::datatypes::TimeUnit::Microsecond;
+use datafusion_expr::Volatility::Immutable;
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkWidthBucket {
+ signature: Signature,
+}
+
+impl Default for SparkWidthBucket {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkWidthBucket {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::user_defined(Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkWidthBucket {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "width_bucket"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(Int32)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ make_scalar_function(width_bucket_kern, vec![])(&args.args)
+ }
+
+ fn coerce_types(&self, types: &[DataType]) -> Result<Vec<DataType>> {
+ use DataType::*;
+
+ let (v, lo, hi, n) = (&types[0], &types[1], &types[2], &types[3]);
+
+ let is_num = |t: &DataType| {
+ matches!(
+ t,
+ Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Decimal128(_, _)
+ )
+ };
+ let is_int = |t: &DataType| matches!(t, Int8 | Int16 | Int32 | Int64);
+
+ if is_num(v) && is_num(lo) && is_num(hi) && is_int(n) {
+ return Ok(vec![Float64, Float64, Float64, Int32]);
+ }
+
+ let all_duration = matches!(v, Duration(_))
+ && matches!(lo, Duration(_))
+ && matches!(hi, Duration(_));
+ if all_duration && is_int(n) {
+ return Ok(vec![
+ Duration(Microsecond),
+ Duration(Microsecond),
+ Duration(Microsecond),
+ Int32,
+ ]);
+ }
+
+ if matches!(v, Interval(YearMonth))
+ && matches!(lo, Interval(YearMonth))
+ && matches!(hi, Interval(YearMonth))
+ && is_int(n)
+ {
+ return Ok(vec![
+ Interval(YearMonth),
+ Interval(YearMonth),
+ Interval(YearMonth),
+ Int32,
+ ]);
+ }
+
+ exec_err!("rint expects a numeric argument, got {}", types[0])
+ }
+
+ fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {
+ // FIXME copy
Review Comment:
this I am not sure if necesary output_ordering, for this I put this, only for these,
If its correct, drop de FIXME, thanks
--
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): implement Spark `width_bucket` function [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #17331:
URL: https://github.com/apache/datafusion/pull/17331#discussion_r2311265828
##########
datafusion/sqllogictest/test_files/spark/math/width_bucket.slt:
##########
@@ -21,42 +21,58 @@
# For more information, please see:
# https://github.com/apache/datafusion/issues/15914
-## Original Query: SELECT width_bucket(-0.9, 5.2, 0.5, 2);
-## PySpark 3.5.5 Result: {'width_bucket(-0.9, 5.2, 0.5, 2)': 3, 'typeof(width_bucket(-0.9, 5.2, 0.5, 2))': 'bigint', 'typeof(-0.9)': 'decimal(1,1)', 'typeof(5.2)': 'decimal(2,1)', 'typeof(0.5)': 'decimal(1,1)', 'typeof(2)': 'int'}
-#query
-#SELECT width_bucket(-0.9::decimal(1,1), 5.2::decimal(2,1), 0.5::decimal(1,1), 2::int);
-
-## Original Query: SELECT width_bucket(-2.1, 1.3, 3.4, 3);
-## PySpark 3.5.5 Result: {'width_bucket(-2.1, 1.3, 3.4, 3)': 0, 'typeof(width_bucket(-2.1, 1.3, 3.4, 3))': 'bigint', 'typeof(-2.1)': 'decimal(2,1)', 'typeof(1.3)': 'decimal(2,1)', 'typeof(3.4)': 'decimal(2,1)', 'typeof(3)': 'int'}
-#query
-#SELECT width_bucket(-2.1::decimal(2,1), 1.3::decimal(2,1), 3.4::decimal(2,1), 3::int);
-
-## Original Query: SELECT width_bucket(5.3, 0.2, 10.6, 5);
-## PySpark 3.5.5 Result: {'width_bucket(5.3, 0.2, 10.6, 5)': 3, 'typeof(width_bucket(5.3, 0.2, 10.6, 5))': 'bigint', 'typeof(5.3)': 'decimal(2,1)', 'typeof(0.2)': 'decimal(1,1)', 'typeof(10.6)': 'decimal(3,1)', 'typeof(5)': 'int'}
-#query
-#SELECT width_bucket(5.3::decimal(2,1), 0.2::decimal(1,1), 10.6::decimal(3,1), 5::int);
-
-## Original Query: SELECT width_bucket(8.1, 0.0, 5.7, 4);
-## PySpark 3.5.5 Result: {'width_bucket(8.1, 0.0, 5.7, 4)': 5, 'typeof(width_bucket(8.1, 0.0, 5.7, 4))': 'bigint', 'typeof(8.1)': 'decimal(2,1)', 'typeof(0.0)': 'decimal(1,1)', 'typeof(5.7)': 'decimal(2,1)', 'typeof(4)': 'int'}
-#query
-#SELECT width_bucket(8.1::decimal(2,1), 0.0::decimal(1,1), 5.7::decimal(2,1), 4::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)": 1, "typeof(width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10))": 'bigint', "typeof(INTERVAL '0' DAY)": 'interval day', "typeof(INTERVAL '10' DAY)": 'interval day', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '0' DAY::interval day, INTERVAL '10' DAY::interval day, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)": 1, "typeof(width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10))": 'bigint', "typeof(INTERVAL '0' YEAR)": 'interval year', "typeof(INTERVAL '10' YEAR)": 'interval year', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '0' YEAR::interval year, INTERVAL '10' YEAR::interval year, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)": 2, "typeof(width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10))": 'bigint', "typeof(INTERVAL '1' DAY)": 'interval day', "typeof(INTERVAL '0' DAY)": 'interval day', "typeof(INTERVAL '10' DAY)": 'interval day', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '1' DAY::interval day, INTERVAL '0' DAY::interval day, INTERVAL '10' DAY::interval day, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)": 2, "typeof(width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10))": 'bigint', "typeof(INTERVAL '1' YEAR)": 'interval year', "typeof(INTERVAL '0' YEAR)": 'interval year', "typeof(INTERVAL '10' YEAR)": 'interval year', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '1' YEAR::interval year, INTERVAL '0' YEAR::interval year, INTERVAL '10' YEAR::interval year, 10::int);
+query I
+SELECT width_bucket(-0.9, 5.2, 0.5, 2)
+----
+3
+
+query I
+SELECT width_bucket(-2.1, 1.3, 3.4, 3)
+----
+0
+
+query I
+SELECT width_bucket(5.3, 0.2, 10.6, 5)
+----
+3
+
+query I
+SELECT width_bucket(8.1, 0.0, 5.7, 4)
+----
+5
+
+query I
+SELECT width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)
+----
+1
+
+query I
+SELECT width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)
+----
+1
+
+query I
+SELECT width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)
+----
+2
+
+query I
+SELECT width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)
+----
+2
+
+# test of sail
+query I
+SELECT width_bucket(0.0, 10.0, 0.0, 5)
+----
+6
+
+query I
+SELECT width_bucket(10.0, 0.0, 10.0, 5)
+----
+6
+
+query I
+SELECT width_bucket(10.0, 0.0, 0.0, 5)
Review Comment:
Can the test with null/nan be added? to use nan please check snippet
```
isnan(-'NaN'::FLOAT)
or
select greatest(arrow_cast('NAN','Float64'), arrow_cast('NAN','Float64'))
```
--
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): implement Spark `width_bucket` function [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #17331:
URL: https://github.com/apache/datafusion/pull/17331#discussion_r2311223397
##########
datafusion/spark/src/function/math/width_bucket.rs:
##########
@@ -0,0 +1,794 @@
+// 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::any::Any;
+use std::sync::Arc;
+
+use crate::function::error_utils::unsupported_data_types_exec_err;
+use arrow::array::{
+ Array, ArrayRef, DurationMicrosecondArray, Float64Array, IntervalMonthDayNanoArray,
+ IntervalYearMonthArray,
+};
+use arrow::datatypes::DataType;
+use arrow::datatypes::DataType::{Duration, Float64, Int32, Interval};
+use arrow::datatypes::IntervalUnit::{MonthDayNano, YearMonth};
+use datafusion_common::cast::{
+ as_duration_microsecond_array, as_float64_array, as_int32_array,
+ as_interval_mdn_array, as_interval_ym_array,
+};
+use datafusion_common::{exec_err, Result};
+use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature};
+use datafusion_functions::utils::make_scalar_function;
+
+use arrow::array::{Int32Array, Int32Builder};
+use arrow::datatypes::TimeUnit::Microsecond;
+use datafusion_expr::Volatility::Immutable;
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkWidthBucket {
+ signature: Signature,
+}
+
+impl Default for SparkWidthBucket {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkWidthBucket {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::user_defined(Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkWidthBucket {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "width_bucket"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(Int32)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ make_scalar_function(width_bucket_kern, vec![])(&args.args)
+ }
+
+ fn coerce_types(&self, types: &[DataType]) -> Result<Vec<DataType>> {
+ use DataType::*;
+
+ let (v, lo, hi, n) = (&types[0], &types[1], &types[2], &types[3]);
+
+ let is_num = |t: &DataType| {
+ matches!(
+ t,
+ Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Decimal128(_, _)
+ )
+ };
+ let is_int = |t: &DataType| matches!(t, Int8 | Int16 | Int32 | Int64);
+
+ if is_num(v) && is_num(lo) && is_num(hi) && is_int(n) {
+ return Ok(vec![Float64, Float64, Float64, Int32]);
+ }
+
+ let all_duration = matches!(v, Duration(_))
Review Comment:
its probably can be rewritten using pattern matching?
--
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