Jefffrey commented on code in PR #18836:
URL: https://github.com/apache/datafusion/pull/18836#discussion_r2544477727
##########
datafusion/spark/src/function/bitwise/bit_get.rs:
##########
@@ -130,164 +108,13 @@ fn spark_bit_get_inner<T: ArrowPrimitiveType>(
Ok(result)
}
-pub fn spark_bit_get(args: &[ArrayRef]) -> Result<ArrayRef> {
- if args.len() != 2 {
- return exec_err!("`bit_get` expects exactly two arguments");
- }
-
- if args[1].data_type() != &Int32 {
- return exec_err!("`bit_get` expects Int32 as the second argument");
- }
-
- let pos_arg = args[1].as_primitive::<Int32Type>();
-
- let ret = match &args[0].data_type() {
- Int64 => {
- let value_arg = args[0].as_primitive::<Int64Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- Int32 => {
- let value_arg = args[0].as_primitive::<Int32Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- Int16 => {
- let value_arg = args[0].as_primitive::<Int16Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- Int8 => {
- let value_arg = args[0].as_primitive::<Int8Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- UInt64 => {
- let value_arg = args[0].as_primitive::<UInt64Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- UInt32 => {
- let value_arg = args[0].as_primitive::<UInt32Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- UInt16 => {
- let value_arg = args[0].as_primitive::<UInt16Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- UInt8 => {
- let value_arg = args[0].as_primitive::<UInt8Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- _ => {
- exec_err!(
- "`bit_get` expects Int64, Int32, Int16, or Int8 as the first
argument"
- )
- }
- }?;
+fn spark_bit_get(args: &[ArrayRef]) -> Result<ArrayRef> {
+ let [value, position] = take_function_args("bit_get", args)?;
+ let pos_arg = position.as_primitive::<Int32Type>();
+ let ret = downcast_integer_array!(
+ value => spark_bit_get_inner(value, pos_arg),
+ DataType::Null => Ok(Int8Array::new_null(value.len())),
+ d => internal_err!("Unsupported datatype for bit_get: {d}"),
+ )?;
Ok(Arc::new(ret))
}
-
-#[cfg(test)]
-mod tests {
Review Comment:
These tests were duplicated by the SLTs so removed them
##########
datafusion/spark/src/function/bitwise/bit_get.rs:
##########
@@ -53,7 +49,17 @@ impl Default for SparkBitGet {
impl SparkBitGet {
pub fn new() -> Self {
Self {
- signature: Signature::user_defined(Volatility::Immutable),
+ signature: Signature::coercible(
Review Comment:
Signature change here
##########
datafusion/spark/src/function/bitwise/bit_get.rs:
##########
@@ -130,164 +108,13 @@ fn spark_bit_get_inner<T: ArrowPrimitiveType>(
Ok(result)
}
-pub fn spark_bit_get(args: &[ArrayRef]) -> Result<ArrayRef> {
- if args.len() != 2 {
- return exec_err!("`bit_get` expects exactly two arguments");
- }
-
- if args[1].data_type() != &Int32 {
- return exec_err!("`bit_get` expects Int32 as the second argument");
- }
-
- let pos_arg = args[1].as_primitive::<Int32Type>();
-
- let ret = match &args[0].data_type() {
- Int64 => {
- let value_arg = args[0].as_primitive::<Int64Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- Int32 => {
- let value_arg = args[0].as_primitive::<Int32Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- Int16 => {
- let value_arg = args[0].as_primitive::<Int16Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- Int8 => {
- let value_arg = args[0].as_primitive::<Int8Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- UInt64 => {
- let value_arg = args[0].as_primitive::<UInt64Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- UInt32 => {
- let value_arg = args[0].as_primitive::<UInt32Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- UInt16 => {
- let value_arg = args[0].as_primitive::<UInt16Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- UInt8 => {
- let value_arg = args[0].as_primitive::<UInt8Type>();
- spark_bit_get_inner(value_arg, pos_arg)
- }
- _ => {
- exec_err!(
- "`bit_get` expects Int64, Int32, Int16, or Int8 as the first
argument"
- )
- }
- }?;
+fn spark_bit_get(args: &[ArrayRef]) -> Result<ArrayRef> {
Review Comment:
This function I made private, and also did a few refactors to make it more
compact.
--
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]
