Skip to content

Commit ccf50f4

Browse files
committed
Allow base64 encoding of fixedsizebinary arrays
1 parent 769f367 commit ccf50f4

File tree

1 file changed

+44
-1
lines changed
  • datafusion/functions/src/encoding

1 file changed

+44
-1
lines changed

datafusion/functions/src/encoding/inner.rs

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,9 @@ use base64::{
2929
Engine as _,
3030
};
3131
use datafusion_common::{
32-
cast::{as_generic_binary_array, as_generic_string_array},
32+
cast::{
33+
as_fixed_size_binary_array, as_generic_binary_array, as_generic_string_array,
34+
},
3335
not_impl_err, plan_err,
3436
utils::take_function_args,
3537
};
@@ -246,6 +248,9 @@ fn encode_process(value: &ColumnarValue, encoding: Encoding) -> Result<ColumnarV
246248
DataType::Utf8View => encoding.encode_utf8_array::<i32>(a.as_ref()),
247249
DataType::Binary => encoding.encode_binary_array::<i32>(a.as_ref()),
248250
DataType::LargeBinary => encoding.encode_binary_array::<i64>(a.as_ref()),
251+
DataType::FixedSizeBinary(_) => {
252+
encoding.encode_fixed_size_binary_array(a.as_ref())
253+
}
249254
other => exec_err!(
250255
"Unsupported data type {other:?} for function encode({encoding})"
251256
),
@@ -265,6 +270,9 @@ fn encode_process(value: &ColumnarValue, encoding: Encoding) -> Result<ColumnarV
265270
),
266271
ScalarValue::LargeBinary(a) => Ok(encoding
267272
.encode_large_scalar(a.as_ref().map(|v: &Vec<u8>| v.as_slice()))),
273+
ScalarValue::FixedSizeBinary(_, a) => Ok(
274+
encoding.encode_scalar(a.as_ref().map(|v: &Vec<u8>| v.as_slice()))
275+
),
268276
other => exec_err!(
269277
"Unsupported data type {other:?} for function encode({encoding})"
270278
),
@@ -401,6 +409,15 @@ impl Encoding {
401409
Ok(ColumnarValue::Array(array))
402410
}
403411

412+
fn encode_fixed_size_binary_array(self, value: &dyn Array) -> Result<ColumnarValue> {
413+
let input_value = as_fixed_size_binary_array(value)?;
414+
let array: ArrayRef = match self {
415+
Self::Base64 => encode_to_array!(base64_encode, input_value),
416+
Self::Hex => encode_to_array!(hex_encode, input_value),
417+
};
418+
Ok(ColumnarValue::Array(array))
419+
}
420+
404421
fn encode_utf8_array<T>(self, value: &dyn Array) -> Result<ColumnarValue>
405422
where
406423
T: OffsetSizeTrait,
@@ -553,3 +570,29 @@ fn decode(args: &[ColumnarValue]) -> Result<ColumnarValue> {
553570
}?;
554571
decode_process(expression, encoding)
555572
}
573+
574+
#[cfg(test)]
575+
mod tests {
576+
#[test]
577+
fn test_encode_fsb() {
578+
use super::*;
579+
580+
let value = vec![0u8; 16];
581+
let array = arrow::array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
582+
vec![Some(value)].into_iter(),
583+
16,
584+
)
585+
.unwrap();
586+
let value = ColumnarValue::Array(Arc::new(array));
587+
588+
let ColumnarValue::Array(result) =
589+
encode_process(&value, Encoding::Base64).unwrap()
590+
else {
591+
panic!("unexpected value");
592+
};
593+
594+
let string_array = result.as_any().downcast_ref::<StringArray>().unwrap();
595+
let result_value = string_array.value(0);
596+
assert_eq!(result_value, "AAAAAAAAAAAAAAAAAAAAAA");
597+
}
598+
}

0 commit comments

Comments
 (0)