This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 9b2cb25c85 arrow-cast: support packing to Dictionary(_,
Utf8View/BinaryView) (#9220)
9b2cb25c85 is described below
commit 9b2cb25c85fab44b5dba7372e0d8f9c7dede3212
Author: Ethan Urbanski <[email protected]>
AuthorDate: Fri Jan 30 09:33:20 2026 -0500
arrow-cast: support packing to Dictionary(_, Utf8View/BinaryView) (#9220)
Fix arrow-cast failing to cast into Dictionary(K, Utf8View) and
Dictionary(K, BinaryView) even though can_cast_types returns true.
# Which issue does this PR close?
- Closes #9219.
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
# Rationale for this change
Downstream pipelines may require view-typed schemas (e.g.
schema_force_view_types) while still dictionary-encoding columns for
memory efficiency. Even with view types, dictionary encoding can reduce
per-row memory (1–4 byte keys vs 16-byte views) for low-cardinality
columns. Today there's no way to materialize Dictionary(K,
Utf8View/BinaryView) via arrow-cast.
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
Add Utf8View and BinaryView support in cast_to_dictionary using a
two-step pack to recast approach:
- Pack to Dictionary(K, Utf8/Binary)
- Cast dictionary values to Utf8View/BinaryView via existing cast
machinery
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
- StringArray/StringViewArray to Dictionary(UInt16, Utf8View) (including
sliced inputs)
- BinaryArray/BinaryViewArray to Dictionary(UInt16, BinaryView)
(including sliced inputs)
- Key overflow for Dictionary(UInt8, Utf8View)
- Empty and all-null inputs
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
# Are there any user-facing changes?
No
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
---
arrow-cast/src/cast/dictionary.rs | 26 +++
arrow-cast/src/cast/mod.rs | 370 +++++++++++++++++++++++++++++++++++++-
2 files changed, 388 insertions(+), 8 deletions(-)
diff --git a/arrow-cast/src/cast/dictionary.rs
b/arrow-cast/src/cast/dictionary.rs
index 64d77236cc..0092bc0c87 100644
--- a/arrow-cast/src/cast/dictionary.rs
+++ b/arrow-cast/src/cast/dictionary.rs
@@ -272,6 +272,19 @@ pub(crate) fn cast_to_dictionary<K:
ArrowDictionaryKeyType>(
}
pack_byte_to_dictionary::<K, GenericStringType<i64>>(array,
cast_options)
}
+ Utf8View => {
+ let base_value_type = match array.data_type() {
+ DataType::LargeUtf8 | DataType::Utf8View =>
DataType::LargeUtf8,
+ _ => DataType::Utf8,
+ };
+
+ let dict_base = cast_to_dictionary::<K>(array, &base_value_type,
cast_options)?;
+ dictionary_cast::<K>(
+ dict_base.as_ref(),
+ &DataType::Dictionary(Box::new(K::DATA_TYPE),
Box::new(DataType::Utf8View)),
+ cast_options,
+ )
+ }
Binary => {
// If the input is a view type, we can avoid casting (thus
copying) the data
if array.data_type() == &DataType::BinaryView {
@@ -286,6 +299,19 @@ pub(crate) fn cast_to_dictionary<K:
ArrowDictionaryKeyType>(
}
pack_byte_to_dictionary::<K, GenericBinaryType<i64>>(array,
cast_options)
}
+ BinaryView => {
+ let base_value_type = match array.data_type() {
+ DataType::LargeBinary | DataType::BinaryView =>
DataType::LargeBinary,
+ _ => DataType::Binary,
+ };
+
+ let dict_base = cast_to_dictionary::<K>(array, &base_value_type,
cast_options)?;
+ dictionary_cast::<K>(
+ dict_base.as_ref(),
+ &DataType::Dictionary(Box::new(K::DATA_TYPE),
Box::new(DataType::BinaryView)),
+ cast_options,
+ )
+ }
FixedSizeBinary(byte_size) => {
pack_byte_to_fixed_size_dictionary::<K>(array, cast_options,
byte_size)
}
diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs
index 6c1629a820..736a427ec0 100644
--- a/arrow-cast/src/cast/mod.rs
+++ b/arrow-cast/src/cast/mod.rs
@@ -5676,10 +5676,7 @@ mod tests {
vec!["4869696969", "48656c6c6f", "4869696969", "null"]
);
// dictionary should only have two distinct values
- let dict_array = cast_array
- .as_any()
- .downcast_ref::<DictionaryArray<Int8Type>>()
- .unwrap();
+ let dict_array = cast_array.as_dictionary::<Int8Type>();
assert_eq!(dict_array.values().len(), 2);
}
@@ -5711,11 +5708,368 @@ mod tests {
]
);
// dictionary should only have three distinct values
- let dict_array = cast_array
- .as_any()
- .downcast_ref::<DictionaryArray<Int8Type>>()
- .unwrap();
+ let dict_array = cast_array.as_dictionary::<Int8Type>();
+ assert_eq!(dict_array.values().len(), 3);
+ }
+
+ #[test]
+ fn test_cast_string_array_to_dict_utf8_view() {
+ let array = StringArray::from(vec![Some("one"), None, Some("three"),
Some("one")]);
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::Utf8View));
+ assert!(can_cast_types(array.data_type(), &cast_type));
+ let cast_array = cast(&array, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+
+ let dict_array = cast_array.as_dictionary::<UInt16Type>();
+ assert_eq!(dict_array.values().data_type(), &DataType::Utf8View);
+ assert_eq!(dict_array.values().len(), 2); // "one" and "three"
deduplicated
+
+ let typed = dict_array.downcast_dict::<StringViewArray>().unwrap();
+ let actual: Vec<Option<&str>> = typed.into_iter().collect();
+ assert_eq!(actual, vec![Some("one"), None, Some("three"),
Some("one")]);
+
+ let keys = dict_array.keys();
+ assert!(keys.is_null(1));
+ assert_eq!(keys.value(0), keys.value(3));
+ assert_ne!(keys.value(0), keys.value(2));
+ }
+
+ #[test]
+ fn test_cast_string_array_to_dict_utf8_view_null_vs_literal_null() {
+ let array = StringArray::from(vec![Some("one"), None, Some("null"),
Some("one")]);
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::Utf8View));
+ assert!(can_cast_types(array.data_type(), &cast_type));
+ let cast_array = cast(&array, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+
+ let dict_array = cast_array.as_dictionary::<UInt16Type>();
+ assert_eq!(dict_array.values().data_type(), &DataType::Utf8View);
+ assert_eq!(dict_array.values().len(), 2);
+
+ let typed = dict_array.downcast_dict::<StringViewArray>().unwrap();
+ let actual: Vec<Option<&str>> = typed.into_iter().collect();
+ assert_eq!(actual, vec![Some("one"), None, Some("null"), Some("one")]);
+
+ let keys = dict_array.keys();
+ assert!(keys.is_null(1));
+ assert_eq!(keys.value(0), keys.value(3));
+ assert_ne!(keys.value(0), keys.value(2));
+ }
+
+ #[test]
+ fn test_cast_string_view_array_to_dict_utf8_view() {
+ let array = StringViewArray::from(vec![Some("one"), None,
Some("three"), Some("one")]);
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::Utf8View));
+ assert!(can_cast_types(array.data_type(), &cast_type));
+ let cast_array = cast(&array, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+
+ let dict_array = cast_array.as_dictionary::<UInt16Type>();
+ assert_eq!(dict_array.values().data_type(), &DataType::Utf8View);
+ assert_eq!(dict_array.values().len(), 2); // "one" and "three"
deduplicated
+
+ let typed = dict_array.downcast_dict::<StringViewArray>().unwrap();
+ let actual: Vec<Option<&str>> = typed.into_iter().collect();
+ assert_eq!(actual, vec![Some("one"), None, Some("three"),
Some("one")]);
+
+ let keys = dict_array.keys();
+ assert!(keys.is_null(1));
+ assert_eq!(keys.value(0), keys.value(3));
+ assert_ne!(keys.value(0), keys.value(2));
+ }
+
+ #[test]
+ fn test_cast_string_view_slice_to_dict_utf8_view() {
+ let array = StringViewArray::from(vec![
+ Some("zero"),
+ Some("one"),
+ None,
+ Some("three"),
+ Some("one"),
+ ]);
+ let view = array.slice(1, 4);
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::Utf8View));
+ assert!(can_cast_types(view.data_type(), &cast_type));
+ let cast_array = cast(&view, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+
+ let dict_array = cast_array.as_dictionary::<UInt16Type>();
+ assert_eq!(dict_array.values().data_type(), &DataType::Utf8View);
+ assert_eq!(dict_array.values().len(), 2);
+
+ let typed = dict_array.downcast_dict::<StringViewArray>().unwrap();
+ let actual: Vec<Option<&str>> = typed.into_iter().collect();
+ assert_eq!(actual, vec![Some("one"), None, Some("three"),
Some("one")]);
+
+ let keys = dict_array.keys();
+ assert!(keys.is_null(1));
+ assert_eq!(keys.value(0), keys.value(3));
+ assert_ne!(keys.value(0), keys.value(2));
+ }
+
+ #[test]
+ fn test_cast_binary_array_to_dict_binary_view() {
+ let mut builder = GenericBinaryBuilder::<i32>::new();
+ builder.append_value(b"hello");
+ builder.append_value(b"hiiii");
+ builder.append_value(b"hiiii"); // duplicate
+ builder.append_null();
+ builder.append_value(b"rustt");
+
+ let array = builder.finish();
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::BinaryView));
+ assert!(can_cast_types(array.data_type(), &cast_type));
+ let cast_array = cast(&array, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+
+ let dict_array = cast_array.as_dictionary::<UInt16Type>();
+ assert_eq!(dict_array.values().data_type(), &DataType::BinaryView);
assert_eq!(dict_array.values().len(), 3);
+
+ let typed = dict_array.downcast_dict::<BinaryViewArray>().unwrap();
+ let actual: Vec<Option<&[u8]>> = typed.into_iter().collect();
+ assert_eq!(
+ actual,
+ vec![
+ Some(b"hello".as_slice()),
+ Some(b"hiiii".as_slice()),
+ Some(b"hiiii".as_slice()),
+ None,
+ Some(b"rustt".as_slice())
+ ]
+ );
+
+ let keys = dict_array.keys();
+ assert!(keys.is_null(3));
+ assert_eq!(keys.value(1), keys.value(2));
+ assert_ne!(keys.value(0), keys.value(1));
+ }
+
+ #[test]
+ fn test_cast_binary_view_array_to_dict_binary_view() {
+ let view = BinaryViewArray::from_iter([
+ Some(b"hello".as_slice()),
+ Some(b"hiiii".as_slice()),
+ Some(b"hiiii".as_slice()), // duplicate
+ None,
+ Some(b"rustt".as_slice()),
+ ]);
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::BinaryView));
+ assert!(can_cast_types(view.data_type(), &cast_type));
+ let cast_array = cast(&view, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+
+ let dict_array = cast_array.as_dictionary::<UInt16Type>();
+ assert_eq!(dict_array.values().data_type(), &DataType::BinaryView);
+ assert_eq!(dict_array.values().len(), 3);
+
+ let typed = dict_array.downcast_dict::<BinaryViewArray>().unwrap();
+ let actual: Vec<Option<&[u8]>> = typed.into_iter().collect();
+ assert_eq!(
+ actual,
+ vec![
+ Some(b"hello".as_slice()),
+ Some(b"hiiii".as_slice()),
+ Some(b"hiiii".as_slice()),
+ None,
+ Some(b"rustt".as_slice())
+ ]
+ );
+
+ let keys = dict_array.keys();
+ assert!(keys.is_null(3));
+ assert_eq!(keys.value(1), keys.value(2));
+ assert_ne!(keys.value(0), keys.value(1));
+ }
+
+ #[test]
+ fn test_cast_binary_view_slice_to_dict_binary_view() {
+ let view = BinaryViewArray::from_iter([
+ Some(b"hello".as_slice()),
+ Some(b"hiiii".as_slice()),
+ Some(b"hiiii".as_slice()), // duplicate
+ None,
+ Some(b"rustt".as_slice()),
+ ]);
+ let sliced = view.slice(1, 4);
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::BinaryView));
+ assert!(can_cast_types(sliced.data_type(), &cast_type));
+ let cast_array = cast(&sliced, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+
+ let dict_array = cast_array.as_dictionary::<UInt16Type>();
+ assert_eq!(dict_array.values().data_type(), &DataType::BinaryView);
+ assert_eq!(dict_array.values().len(), 2);
+
+ let typed = dict_array.downcast_dict::<BinaryViewArray>().unwrap();
+ let actual: Vec<Option<&[u8]>> = typed.into_iter().collect();
+ assert_eq!(
+ actual,
+ vec![
+ Some(b"hiiii".as_slice()),
+ Some(b"hiiii".as_slice()),
+ None,
+ Some(b"rustt".as_slice())
+ ]
+ );
+
+ let keys = dict_array.keys();
+ assert!(keys.is_null(2));
+ assert_eq!(keys.value(0), keys.value(1));
+ assert_ne!(keys.value(0), keys.value(3));
+ }
+
+ #[test]
+ fn test_cast_string_array_to_dict_utf8_view_key_overflow_u8() {
+ let array = StringArray::from_iter_values((0..257).map(|i|
format!("v{i}")));
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt8),
Box::new(DataType::Utf8View));
+ assert!(can_cast_types(array.data_type(), &cast_type));
+ let err = cast(&array, &cast_type).unwrap_err();
+ assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
+ }
+
+ #[test]
+ fn test_cast_large_string_array_to_dict_utf8_view() {
+ let array = LargeStringArray::from(vec![Some("one"), None,
Some("three"), Some("one")]);
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::Utf8View));
+ assert!(can_cast_types(array.data_type(), &cast_type));
+ let cast_array = cast(&array, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+
+ let dict_array = cast_array.as_dictionary::<UInt16Type>();
+ assert_eq!(dict_array.values().data_type(), &DataType::Utf8View);
+ assert_eq!(dict_array.values().len(), 2); // "one" and "three"
deduplicated
+
+ let typed = dict_array.downcast_dict::<StringViewArray>().unwrap();
+ let actual: Vec<Option<&str>> = typed.into_iter().collect();
+ assert_eq!(actual, vec![Some("one"), None, Some("three"),
Some("one")]);
+
+ let keys = dict_array.keys();
+ assert!(keys.is_null(1));
+ assert_eq!(keys.value(0), keys.value(3));
+ assert_ne!(keys.value(0), keys.value(2));
+ }
+
+ #[test]
+ fn test_cast_large_binary_array_to_dict_binary_view() {
+ let mut builder = GenericBinaryBuilder::<i64>::new();
+ builder.append_value(b"hello");
+ builder.append_value(b"world");
+ builder.append_value(b"hello"); // duplicate
+ builder.append_null();
+
+ let array = builder.finish();
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::BinaryView));
+ assert!(can_cast_types(array.data_type(), &cast_type));
+ let cast_array = cast(&array, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+
+ let dict_array = cast_array.as_dictionary::<UInt16Type>();
+ assert_eq!(dict_array.values().data_type(), &DataType::BinaryView);
+ assert_eq!(dict_array.values().len(), 2); // "hello" and "world"
deduplicated
+
+ let typed = dict_array.downcast_dict::<BinaryViewArray>().unwrap();
+ let actual: Vec<Option<&[u8]>> = typed.into_iter().collect();
+ assert_eq!(
+ actual,
+ vec![
+ Some(b"hello".as_slice()),
+ Some(b"world".as_slice()),
+ Some(b"hello".as_slice()),
+ None
+ ]
+ );
+
+ let keys = dict_array.keys();
+ assert!(keys.is_null(3));
+ assert_eq!(keys.value(0), keys.value(2));
+ assert_ne!(keys.value(0), keys.value(1));
+ }
+
+ #[test]
+ fn test_cast_empty_string_array_to_dict_utf8_view() {
+ let array = StringArray::from(Vec::<Option<&str>>::new());
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::Utf8View));
+ assert!(can_cast_types(array.data_type(), &cast_type));
+ let cast_array = cast(&array, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+ assert_eq!(cast_array.len(), 0);
+ }
+
+ #[test]
+ fn test_cast_empty_binary_array_to_dict_binary_view() {
+ let array = BinaryArray::from(Vec::<Option<&[u8]>>::new());
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::BinaryView));
+ assert!(can_cast_types(array.data_type(), &cast_type));
+ let cast_array = cast(&array, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+ assert_eq!(cast_array.len(), 0);
+ }
+
+ #[test]
+ fn test_cast_all_null_string_array_to_dict_utf8_view() {
+ let array = StringArray::from(vec![None::<&str>, None, None]);
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::Utf8View));
+ assert!(can_cast_types(array.data_type(), &cast_type));
+ let cast_array = cast(&array, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+ assert_eq!(cast_array.null_count(), 3);
+
+ let dict_array = cast_array.as_dictionary::<UInt16Type>();
+ assert_eq!(dict_array.values().data_type(), &DataType::Utf8View);
+ assert_eq!(dict_array.values().len(), 0);
+ assert_eq!(dict_array.keys().null_count(), 3);
+
+ let typed = dict_array.downcast_dict::<StringViewArray>().unwrap();
+ let actual: Vec<Option<&str>> = typed.into_iter().collect();
+ assert_eq!(actual, vec![None, None, None]);
+ }
+
+ #[test]
+ fn test_cast_all_null_binary_array_to_dict_binary_view() {
+ let array = BinaryArray::from(vec![None::<&[u8]>, None, None]);
+
+ let cast_type =
+ DataType::Dictionary(Box::new(DataType::UInt16),
Box::new(DataType::BinaryView));
+ assert!(can_cast_types(array.data_type(), &cast_type));
+ let cast_array = cast(&array, &cast_type).unwrap();
+ assert_eq!(cast_array.data_type(), &cast_type);
+ assert_eq!(cast_array.null_count(), 3);
+
+ let dict_array = cast_array.as_dictionary::<UInt16Type>();
+ assert_eq!(dict_array.values().data_type(), &DataType::BinaryView);
+ assert_eq!(dict_array.values().len(), 0);
+ assert_eq!(dict_array.keys().null_count(), 3);
+
+ let typed = dict_array.downcast_dict::<BinaryViewArray>().unwrap();
+ let actual: Vec<Option<&[u8]>> = typed.into_iter().collect();
+ assert_eq!(actual, vec![None, None, None]);
}
#[test]
