You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by tu...@apache.org on 2023/01/16 17:33:52 UTC
[arrow-rs] branch master updated: Support casting from binary to dictionary of binary (#3482)
This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 07fd43494 Support casting from binary to dictionary of binary (#3482)
07fd43494 is described below
commit 07fd434946e262ec70b303b5065fd75fc3bde922
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Mon Jan 16 09:33:47 2023 -0800
Support casting from binary to dictionary of binary (#3482)
---
.../builder/generic_bytes_dictionary_builder.rs | 4 ++--
arrow-cast/src/cast.rs | 26 ++++++++++++++++++++++
arrow/tests/array_cast.rs | 2 ++
3 files changed, 30 insertions(+), 2 deletions(-)
diff --git a/arrow-array/src/builder/generic_bytes_dictionary_builder.rs b/arrow-array/src/builder/generic_bytes_dictionary_builder.rs
index 34b736d65..4a920f3ee 100644
--- a/arrow-array/src/builder/generic_bytes_dictionary_builder.rs
+++ b/arrow-array/src/builder/generic_bytes_dictionary_builder.rs
@@ -268,7 +268,7 @@ where
let keys = self.keys_builder.finish();
let data_type =
- DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(DataType::Utf8));
+ DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(T::DATA_TYPE));
let builder = keys
.into_data()
@@ -285,7 +285,7 @@ where
let keys = self.keys_builder.finish_cloned();
let data_type =
- DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(DataType::Utf8));
+ DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(T::DATA_TYPE));
let builder = keys
.into_data()
diff --git a/arrow-cast/src/cast.rs b/arrow-cast/src/cast.rs
index 8b8244a7c..aa6697a71 100644
--- a/arrow-cast/src/cast.rs
+++ b/arrow-cast/src/cast.rs
@@ -3303,6 +3303,8 @@ fn cast_to_dictionary<K: ArrowDictionaryKeyType>(
),
Utf8 => pack_string_to_dictionary::<K>(array, cast_options),
LargeUtf8 => pack_string_to_dictionary::<K>(array, cast_options),
+ Binary => pack_binary_to_dictionary::<K>(array, cast_options),
+ LargeBinary => pack_binary_to_dictionary::<K>(array, cast_options),
_ => Err(ArrowError::CastError(format!(
"Unsupported output type for dictionary packing: {:?}",
dict_value_type
@@ -3366,6 +3368,30 @@ where
Ok(Arc::new(b.finish()))
}
+// Packs the data as a BinaryDictionaryArray, if possible, with the
+// key types of K
+fn pack_binary_to_dictionary<K>(
+ array: &ArrayRef,
+ cast_options: &CastOptions,
+) -> Result<ArrayRef, ArrowError>
+where
+ K: ArrowDictionaryKeyType,
+{
+ let cast_values = cast_with_options(array, &DataType::Binary, cast_options)?;
+ let values = cast_values.as_any().downcast_ref::<BinaryArray>().unwrap();
+ let mut b = BinaryDictionaryBuilder::<K>::with_capacity(values.len(), 1024, 1024);
+
+ // copy each element one at a time
+ for i in 0..values.len() {
+ if values.is_null(i) {
+ b.append_null();
+ } else {
+ b.append(values.value(i))?;
+ }
+ }
+ Ok(Arc::new(b.finish()))
+}
+
/// Helper function that takes a primitive array and casts to a (generic) list array.
fn cast_primitive_to_list<OffsetSize: OffsetSizeTrait + NumCast>(
array: &ArrayRef,
diff --git a/arrow/tests/array_cast.rs b/arrow/tests/array_cast.rs
index be37a7636..91d2da998 100644
--- a/arrow/tests/array_cast.rs
+++ b/arrow/tests/array_cast.rs
@@ -411,7 +411,9 @@ fn get_all_types() -> Vec<DataType> {
),
Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)),
+ Dictionary(Box::new(DataType::Int16), Box::new(DataType::Binary)),
Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
+ Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Binary)),
Decimal128(38, 0),
Dictionary(Box::new(DataType::Int8), Box::new(Decimal128(38, 0))),
Dictionary(Box::new(DataType::Int16), Box::new(Decimal128(38, 0))),