Skip to content

Commit 35c1cfd

Browse files
authored
Introduce TypeSignatureClass::Binary to allow accepting arbitrarily sized FixedSizeBinary arguments (#17531)
* Introduce wildcard const for FixedSizeBinary type signature * Add Binary to TypeSignatureClass * Remove FIXED_SIZE_BINARY_WILDCARD
1 parent 4a5b137 commit 35c1cfd

File tree

3 files changed

+23
-13
lines changed

3 files changed

+23
-13
lines changed

datafusion/common/src/types/native.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -472,4 +472,9 @@ impl NativeType {
472472
pub fn is_duration(&self) -> bool {
473473
matches!(self, NativeType::Duration(_))
474474
}
475+
476+
#[inline]
477+
pub fn is_binary(&self) -> bool {
478+
matches!(self, NativeType::Binary | NativeType::FixedSizeBinary(_))
479+
}
475480
}

datafusion/expr-common/src/signature.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,8 @@ pub enum TypeSignatureClass {
266266
// TODO:
267267
// Numeric
268268
Integer,
269+
/// Encompasses both the native Binary as well as arbitrarily sized FixedSizeBinary types
270+
Binary,
269271
}
270272

271273
impl Display for TypeSignatureClass {
@@ -303,6 +305,9 @@ impl TypeSignatureClass {
303305
TypeSignatureClass::Integer => {
304306
vec![DataType::Int64]
305307
}
308+
TypeSignatureClass::Binary => {
309+
vec![DataType::Binary]
310+
}
306311
}
307312
}
308313

@@ -322,6 +327,7 @@ impl TypeSignatureClass {
322327
TypeSignatureClass::Interval if logical_type.is_interval() => true,
323328
TypeSignatureClass::Duration if logical_type.is_duration() => true,
324329
TypeSignatureClass::Integer if logical_type.is_integer() => true,
330+
TypeSignatureClass::Binary if logical_type.is_binary() => true,
325331
_ => false,
326332
}
327333
}
@@ -352,6 +358,9 @@ impl TypeSignatureClass {
352358
TypeSignatureClass::Integer if native_type.is_integer() => {
353359
Ok(origin_type.to_owned())
354360
}
361+
TypeSignatureClass::Binary if native_type.is_binary() => {
362+
Ok(origin_type.to_owned())
363+
}
355364
_ => internal_err!("May miss the matching logic in `matches_native_type`"),
356365
}
357366
}

datafusion/spark/src/function/bitmap/bitmap_count.rs

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,10 @@ use arrow::datatypes::DataType::{
2727
Binary, BinaryView, FixedSizeBinary, Int64, LargeBinary,
2828
};
2929
use datafusion_common::utils::take_function_args;
30-
use datafusion_common::{internal_datafusion_err, internal_err, plan_err, Result};
30+
use datafusion_common::{internal_datafusion_err, internal_err, Result};
3131
use datafusion_expr::{
32-
ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
32+
Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature,
33+
TypeSignatureClass, Volatility,
3334
};
3435
use datafusion_functions::utils::make_scalar_function;
3536
use datafusion_functions::{downcast_arg, downcast_named_arg};
@@ -48,8 +49,10 @@ impl Default for BitmapCount {
4849
impl BitmapCount {
4950
pub fn new() -> Self {
5051
Self {
51-
// TODO: add definitive TypeSignature after https://github.com/apache/datafusion/issues/17291 is done
52-
signature: Signature::any(1, Volatility::Immutable),
52+
signature: Signature::coercible(
53+
vec![Coercion::new_exact(TypeSignatureClass::Binary)],
54+
Volatility::Immutable,
55+
),
5356
}
5457
}
5558
}
@@ -67,15 +70,8 @@ impl ScalarUDFImpl for BitmapCount {
6770
&self.signature
6871
}
6972

70-
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
71-
match arg_types.first() {
72-
Some(Binary | BinaryView | FixedSizeBinary(_) | LargeBinary) => Ok(Int64),
73-
Some(data_type) => plan_err!(
74-
"bitmap_count expects Binary/BinaryView/FixedSizeBinary/LargeBinary as argument, got {:?}",
75-
data_type
76-
),
77-
None => internal_err!("bitmap_count does not support zero arguments"),
78-
}
73+
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
74+
Ok(Int64)
7975
}
8076

8177
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {

0 commit comments

Comments
 (0)