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 c258274aef test(parquet): move nested writer roundtrips —— Part II
(#11183)
c258274aef is described below
commit c258274aef28810584f1ef356fe145348a075d59
Author: Phoenix <[email protected]>
AuthorDate: Sat Sep 26 00:53:27 2026 +0800
test(parquet): move nested writer roundtrips —— Part II (#11183)
# Which issue does this PR close?
Part of apache/arrow-rs#10540.
# Stack base
This draft PR targets `apache/arrow-rs:main`. It follows the first
stage, which was merged in apache/arrow-rs#11128. The branch now
contains only the three stage-II commits on top of `main`, so this PR's
diff is the stage-II change.
# Rationale for this change
Keep the next group of writer roundtrip tests together in the
integration suite, following the shared harness and scalar tests in
#11128.
# What changes are included in this PR?
Move nested list/map/struct and dictionary roundtrip tests from the
writer unit-test module to `parquet/tests/arrow_writer/roundtrip.rs`.
The three commits preserve the existing assertions and test
configurations.
# Are these changes tested?
- `cargo test --locked -p parquet --test arrow_writer` — 88 passed.
- `cargo test --locked -p parquet --lib -- arrow::arrow_writer::tests` —
80 passed.
- `cargo fmt --all -- --check`
- `git diff --check upstream/main..HEAD`
# Are there any user-facing changes?
No. This is a test-only move.
---
parquet/src/arrow/arrow_writer/mod.rs | 921 +-----------------------------
parquet/tests/arrow_writer/roundtrip.rs | 972 +++++++++++++++++++++++++++++++-
2 files changed, 960 insertions(+), 933 deletions(-)
diff --git a/parquet/src/arrow/arrow_writer/mod.rs
b/parquet/src/arrow/arrow_writer/mod.rs
index 0c61143948..6b5f29d0f1 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -2103,16 +2103,14 @@ mod tests {
use crate::file::reader::SerializedPageReader;
use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
use crate::schema::types::ColumnPath;
- use arrow::datatypes::ToByteSlice;
use arrow::datatypes::{DataType, Schema};
use arrow::error::Result as ArrowResult;
use arrow::util::data_gen::create_random_array;
use arrow::util::pretty::pretty_format_batches;
use arrow::{array::*, buffer::Buffer};
- use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, NullBuffer,
OffsetBuffer, i256};
+ use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, NullBuffer,
OffsetBuffer};
use arrow_schema::Fields;
use half::f16;
- use num_traits::{FromPrimitive, ToPrimitive};
use tempfile::tempfile;
use crate::basic::{Encoding, EncodingMask};
@@ -2424,218 +2422,6 @@ mod tests {
.run();
}
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_list() {
- // define schema
- let schema = Schema::new(vec![Field::new(
- "a",
- DataType::List(Arc::new(Field::new_list_field(DataType::Int32,
false))),
- true,
- )]);
-
- // create some data
- let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
-
- // Construct a buffer for value offsets, for the nested array:
- // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
- let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6,
10].to_byte_slice());
-
- // Construct a list array from the above two
- let a_list_data =
ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
- DataType::Int32,
- false,
- ))))
- .len(5)
- .add_buffer(a_value_offsets)
- .add_child_data(a_values.into_data())
- .null_bit_buffer(Some(Buffer::from([0b00011011])))
- .build()
- .unwrap();
- let a = ListArray::from(a_list_data);
- assert_eq!(a.null_count(), 1);
-
- RoundTripTest::new(Arc::new(a))
- .with_schema(Arc::new(schema))
- .run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_list_non_null() {
- // define schema
- let schema = Schema::new(vec![Field::new(
- "a",
- DataType::List(Arc::new(Field::new_list_field(DataType::Int32,
false))),
- false,
- )]);
-
- // create some data
- let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
-
- // Construct a buffer for value offsets, for the nested array:
- // [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
- let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6,
10].to_byte_slice());
-
- // Construct a list array from the above two
- let a_list_data =
ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
- DataType::Int32,
- false,
- ))))
- .len(5)
- .add_buffer(a_value_offsets)
- .add_child_data(a_values.into_data())
- .build()
- .unwrap();
- let a = ListArray::from(a_list_data);
- assert_eq!(a.null_count(), 0);
-
- RoundTripTest::new(Arc::new(a))
- .with_schema(Arc::new(schema))
- .run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_list_view() {
- let list_field = Arc::new(Field::new_list_field(DataType::Int32,
false));
- let schema = Schema::new(vec![Field::new(
- "a",
- DataType::ListView(list_field.clone()),
- true,
- )]);
-
- // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
- let a = ListViewArray::new(
- list_field,
- vec![0, 1, 0, 3, 6].into(),
- vec![1, 2, 0, 3, 4].into(),
- Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
- Some(vec![true, true, false, true, true].into()),
- );
- assert_eq!(a.null_count(), 1);
-
- RoundTripTest::new(Arc::new(a))
- .with_schema(Arc::new(schema))
- .run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_list_view_non_null() {
- let list_field = Arc::new(Field::new_list_field(DataType::Int32,
false));
- let schema = Schema::new(vec![Field::new(
- "a",
- DataType::ListView(list_field.clone()),
- false,
- )]);
-
- // [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
- let a = ListViewArray::new(
- list_field,
- vec![0, 1, 0, 3, 6].into(),
- vec![1, 2, 0, 3, 4].into(),
- Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
- None,
- );
- assert_eq!(a.null_count(), 0);
-
- RoundTripTest::new(Arc::new(a))
- .with_schema(Arc::new(schema))
- .run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_list_view_out_of_order() {
- let list_field = Arc::new(Field::new_list_field(DataType::Int32,
false));
- let schema = Schema::new(vec![Field::new(
- "a",
- DataType::ListView(list_field.clone()),
- false,
- )]);
-
- // [[1], [2, 3], [], [7, 8, 9, 10], [4, 5, 6]] - out of order offsets
- let a = ListViewArray::new(
- list_field,
- vec![0, 1, 0, 6, 3].into(),
- vec![1, 2, 0, 4, 3].into(),
- Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
- None,
- );
- assert_eq!(a.null_count(), 0);
-
- RoundTripTest::new(Arc::new(a))
- .with_schema(Arc::new(schema))
- .run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_large_list_view() {
- let list_field = Arc::new(Field::new_list_field(DataType::Int32,
false));
- let schema = Schema::new(vec![Field::new(
- "a",
- DataType::LargeListView(list_field.clone()),
- true,
- )]);
-
- // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
- let a = LargeListViewArray::new(
- list_field,
- vec![0i64, 1, 0, 3, 6].into(),
- vec![1i64, 2, 0, 3, 4].into(),
- Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
- Some(vec![true, true, false, true, true].into()),
- );
- assert_eq!(a.null_count(), 1);
-
- RoundTripTest::new(Arc::new(a))
- .with_schema(Arc::new(schema))
- .run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_list_view_with_struct() {
- // Test ListView containing Struct: ListView<Struct<Int32, Utf8>>
- let struct_fields = Fields::from(vec![
- Field::new("id", DataType::Int32, false),
- Field::new("name", DataType::Utf8, false),
- ]);
- let struct_type = DataType::Struct(struct_fields.clone());
- let list_field = Arc::new(Field::new("item", struct_type.clone(),
false));
-
- let schema = Schema::new(vec![Field::new(
- "a",
- DataType::ListView(list_field.clone()),
- true,
- )]);
-
- // Create struct values
- let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let name_array = StringArray::from(vec!["a", "b", "c", "d", "e"]);
- let struct_array = StructArray::new(
- struct_fields,
- vec![Arc::new(id_array), Arc::new(name_array)],
- None,
- );
-
- // Create ListView: [{1, "a"}, {2, "b"}], null, [{3, "c"}, {4, "d"},
{5, "e"}]
- let list_view = ListViewArray::new(
- list_field,
- vec![0, 2, 2].into(), // offsets
- vec![2, 0, 3].into(), // sizes
- Arc::new(struct_array),
- Some(vec![true, false, true].into()),
- );
- assert_eq!(list_view.null_count(), 1);
-
- RoundTripTest::new(Arc::new(list_view))
- .with_schema(Arc::new(schema))
- .run();
- }
-
#[test]
#[cfg_attr(miri, ignore)] // Takes too long
fn arrow_writer_binary() {
@@ -2704,32 +2490,6 @@ mod tests {
.run();
}
- /// Test round-trip of Dictionary<UInt32, Utf8View> and
- /// Dictionary<UInt32, BinaryView> typed columns.
- #[test]
- fn arrow_writer_string_view_dictionary() {
- let raw_string_values = vec!["a", "b", "large payload over 12 bytes"];
- let raw_binary_values = vec![
- b"a".to_vec(),
- b"b".to_vec(),
- b"large payload over 12 bytes".to_vec(),
- ];
-
- let keys = UInt32Array::from(vec![Some(0), None, Some(2), Some(1),
None]);
-
- let string_view_values =
Arc::new(StringViewArray::from(raw_string_values));
- let string_dict: ArrayRef = Arc::new(
- DictionaryArray::<UInt32Type>::try_new(keys.clone(),
string_view_values).unwrap(),
- );
-
- let binary_view_values =
Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
- let binary_dict: ArrayRef =
- Arc::new(DictionaryArray::<UInt32Type>::try_new(keys,
binary_view_values).unwrap());
-
- RoundTripTest::new(string_dict).run();
- RoundTripTest::new(binary_dict).run();
- }
-
fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch {
let decimal_field = Field::new("a", DataType::Decimal128(precision,
scale), false);
let schema = Schema::new(vec![decimal_field]);
@@ -2757,390 +2517,6 @@ mod tests {
roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
}
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_complex() {
- // define schema
- let struct_field_d = Arc::new(Field::new("d", DataType::Float64,
true));
- let struct_field_f = Arc::new(Field::new("f", DataType::Float32,
true));
- let struct_field_g = Arc::new(Field::new_list(
- "g",
- Field::new_list_field(DataType::Int16, true),
- false,
- ));
- let struct_field_h = Arc::new(Field::new_list(
- "h",
- Field::new_list_field(DataType::Int16, false),
- true,
- ));
- let struct_field_e = Arc::new(Field::new_struct(
- "e",
- vec![
- struct_field_f.clone(),
- struct_field_g.clone(),
- struct_field_h.clone(),
- ],
- false,
- ));
- let schema = Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("b", DataType::Int32, true),
- Field::new_struct(
- "c",
- vec![struct_field_d.clone(), struct_field_e.clone()],
- false,
- ),
- ]);
-
- // create some data
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
- let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
- let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None,
Some(5.25)]);
-
- let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
-
- // Construct a buffer for value offsets, for the nested array:
- // [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
- let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6,
10].to_byte_slice());
-
- // Construct a list array from the above two
- let g_list_data =
ArrayData::builder(struct_field_g.data_type().clone())
- .len(5)
- .add_buffer(g_value_offsets.clone())
- .add_child_data(g_value.to_data())
- .build()
- .unwrap();
- let g = ListArray::from(g_list_data);
- // The difference between g and h is that h has a null bitmap
- let h_list_data =
ArrayData::builder(struct_field_h.data_type().clone())
- .len(5)
- .add_buffer(g_value_offsets)
- .add_child_data(g_value.to_data())
- .null_bit_buffer(Some(Buffer::from([0b00011011])))
- .build()
- .unwrap();
- let h = ListArray::from(h_list_data);
-
- let e = StructArray::from(vec![
- (struct_field_f, Arc::new(f) as ArrayRef),
- (struct_field_g, Arc::new(g) as ArrayRef),
- (struct_field_h, Arc::new(h) as ArrayRef),
- ]);
-
- let c = StructArray::from(vec![
- (struct_field_d, Arc::new(d) as ArrayRef),
- (struct_field_e, Arc::new(e) as ArrayRef),
- ]);
-
- // build a record batch
- let batch = RecordBatch::try_new(
- Arc::new(schema),
- vec![Arc::new(a), Arc::new(b), Arc::new(c)],
- )
- .unwrap();
-
- roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
- roundtrip(batch, Some(SMALL_SIZE / 3));
- }
-
- #[test]
- fn arrow_writer_complex_mixed() {
- // This test was added while investigating
https://github.com/apache/arrow-rs/issues/244.
- // It was subsequently fixed while investigating
https://github.com/apache/arrow-rs/issues/245.
-
- // define schema
- let offset_field = Arc::new(Field::new("offset", DataType::Int32,
false));
- let partition_field = Arc::new(Field::new("partition",
DataType::Int64, true));
- let topic_field = Arc::new(Field::new("topic", DataType::Utf8, true));
- let schema = Schema::new(vec![Field::new(
- "some_nested_object",
- DataType::Struct(Fields::from(vec![
- offset_field.clone(),
- partition_field.clone(),
- topic_field.clone(),
- ])),
- false,
- )]);
-
- // create some data
- let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let partition = Int64Array::from(vec![Some(1), None, None, Some(4),
Some(5)]);
- let topic = StringArray::from(vec![Some("A"), None, Some("A"),
Some(""), None]);
-
- let some_nested_object = StructArray::from(vec![
- (offset_field, Arc::new(offset) as ArrayRef),
- (partition_field, Arc::new(partition) as ArrayRef),
- (topic_field, Arc::new(topic) as ArrayRef),
- ]);
-
- // build a record batch
- let batch =
- RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(some_nested_object)]).unwrap();
-
- roundtrip(batch, Some(SMALL_SIZE / 2));
- }
-
- #[test]
- fn arrow_writer_map() {
- // Note: we are using the JSON Arrow reader for brevity
- let json_content = r#"
- {"stocks":{"long": "$AAA", "short": "$BBB"}}
- {"stocks":{"long": null, "long": "$CCC", "short": null}}
- {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}}
- "#;
- let entries_struct_type = DataType::Struct(Fields::from(vec![
- Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8,
false),
- Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8,
true),
- ]));
- let stocks_field = Field::new(
- "stocks",
- DataType::Map(
- Arc::new(Field::new(
- Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
- entries_struct_type,
- false,
- )),
- false,
- ),
- true,
- );
- let schema = Arc::new(Schema::new(vec![stocks_field]));
- let builder =
arrow::json::ReaderBuilder::new(schema).with_batch_size(64);
- let mut reader =
builder.build(std::io::Cursor::new(json_content)).unwrap();
-
- let batch = reader.next().unwrap().unwrap();
- roundtrip(batch, None);
- }
-
- #[test]
- fn arrow_writer_2_level_struct() {
- // tests writing <struct<struct<primitive>>
- let field_c = Field::new("c", DataType::Int32, true);
- let field_b = Field::new("b", DataType::Struct(vec![field_c].into()),
true);
- let type_a = DataType::Struct(vec![field_b.clone()].into());
- let field_a = Field::new("a", type_a, true);
- let schema = Schema::new(vec![field_a.clone()]);
-
- // create data
- let c = Int32Array::from(vec![Some(1), None, Some(3), None, None,
Some(6)]);
- let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
- .len(6)
- .null_bit_buffer(Some(Buffer::from([0b00100111])))
- .add_child_data(c.into_data())
- .build()
- .unwrap();
- let b = StructArray::from(b_data);
- let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
- .len(6)
- .null_bit_buffer(Some(Buffer::from([0b00101111])))
- .add_child_data(b.into_data())
- .build()
- .unwrap();
- let a = StructArray::from(a_data);
-
- assert_eq!(a.null_count(), 1);
- assert_eq!(a.column(0).null_count(), 2);
-
- // build a racord batch
- let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(a)]).unwrap();
-
- roundtrip(batch, Some(SMALL_SIZE / 2));
- }
-
- #[test]
- fn arrow_writer_2_level_struct_non_null() {
- // tests writing <struct<struct<primitive>>
- let field_c = Field::new("c", DataType::Int32, false);
- let type_b = DataType::Struct(vec![field_c].into());
- let field_b = Field::new("b", type_b.clone(), false);
- let type_a = DataType::Struct(vec![field_b].into());
- let field_a = Field::new("a", type_a.clone(), false);
- let schema = Schema::new(vec![field_a]);
-
- // create data
- let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
- let b_data = ArrayDataBuilder::new(type_b)
- .len(6)
- .add_child_data(c.into_data())
- .build()
- .unwrap();
- let b = StructArray::from(b_data);
- let a_data = ArrayDataBuilder::new(type_a)
- .len(6)
- .add_child_data(b.into_data())
- .build()
- .unwrap();
- let a = StructArray::from(a_data);
-
- assert_eq!(a.null_count(), 0);
- assert_eq!(a.column(0).null_count(), 0);
-
- // build a racord batch
- let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(a)]).unwrap();
-
- roundtrip(batch, Some(SMALL_SIZE / 2));
- }
-
- #[test]
- fn arrow_writer_2_level_struct_mixed_null() {
- // tests writing <struct<struct<primitive>>
- let field_c = Field::new("c", DataType::Int32, false);
- let type_b = DataType::Struct(vec![field_c].into());
- let field_b = Field::new("b", type_b.clone(), true);
- let type_a = DataType::Struct(vec![field_b].into());
- let field_a = Field::new("a", type_a.clone(), false);
- let schema = Schema::new(vec![field_a]);
-
- // create data
- let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
- let b_data = ArrayDataBuilder::new(type_b)
- .len(6)
- .null_bit_buffer(Some(Buffer::from([0b00100111])))
- .add_child_data(c.into_data())
- .build()
- .unwrap();
- let b = StructArray::from(b_data);
- // a intentionally has no null buffer, to test that this is handled
correctly
- let a_data = ArrayDataBuilder::new(type_a)
- .len(6)
- .add_child_data(b.into_data())
- .build()
- .unwrap();
- let a = StructArray::from(a_data);
-
- assert_eq!(a.null_count(), 0);
- assert_eq!(a.column(0).null_count(), 2);
-
- // build a racord batch
- let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(a)]).unwrap();
-
- roundtrip(batch, Some(SMALL_SIZE / 2));
- }
-
- #[test]
- fn arrow_writer_2_level_struct_mixed_null_2() {
- // tests writing <struct<struct<primitive>>, where the primitive
columns are non-null.
- let field_c = Field::new("c", DataType::Int32, false);
- let field_d = Field::new("d", DataType::FixedSizeBinary(4), false);
- let field_e = Field::new(
- "e",
- DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8)),
- false,
- );
-
- let field_b = Field::new(
- "b",
- DataType::Struct(vec![field_c, field_d, field_e].into()),
- false,
- );
- let type_a = DataType::Struct(vec![field_b.clone()].into());
- let field_a = Field::new("a", type_a, true);
- let schema = Schema::new(vec![field_a.clone()]);
-
- // create data
- let c = Int32Array::from_iter_values(0..6);
- let d = FixedSizeBinaryArray::try_from_iter(
- ["aaaa", "bbbb", "cccc", "dddd", "eeee", "ffff"].into_iter(),
- )
- .expect("four byte values");
- let e = Int32DictionaryArray::from_iter(["one", "two", "three",
"four", "five", "one"]);
- let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
- .len(6)
- .add_child_data(c.into_data())
- .add_child_data(d.into_data())
- .add_child_data(e.into_data())
- .build()
- .unwrap();
- let b = StructArray::from(b_data);
- let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
- .len(6)
- .null_bit_buffer(Some(Buffer::from([0b00100101])))
- .add_child_data(b.into_data())
- .build()
- .unwrap();
- let a = StructArray::from(a_data);
-
- assert_eq!(a.null_count(), 3);
- assert_eq!(a.column(0).null_count(), 0);
-
- // build a record batch
- let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(a)]).unwrap();
-
- roundtrip(batch, Some(SMALL_SIZE / 2));
- }
-
- #[test]
- fn test_fixed_size_binary_in_dict() {
- fn test_fixed_size_binary_in_dict_inner<K>()
- where
- K: ArrowDictionaryKeyType,
- K::Native: FromPrimitive + ToPrimitive + TryFrom<u8>,
- <<K as arrow_array::ArrowPrimitiveType>::Native as
TryFrom<u8>>::Error: std::fmt::Debug,
- {
- let field = Field::new(
- "a",
- DataType::Dictionary(
- Box::new(K::DATA_TYPE),
- Box::new(DataType::FixedSizeBinary(4)),
- ),
- false,
- );
- let schema = Schema::new(vec![field]);
-
- let keys: Vec<K::Native> = vec![
- K::Native::try_from(0u8).unwrap(),
- K::Native::try_from(0u8).unwrap(),
- K::Native::try_from(1u8).unwrap(),
- ];
- let keys = PrimitiveArray::<K>::from_iter_values(keys);
- let values = FixedSizeBinaryArray::try_from_iter(
- vec![vec![0, 0, 0, 0], vec![1, 1, 1, 1]].into_iter(),
- )
- .unwrap();
-
- let data = DictionaryArray::<K>::new(keys, Arc::new(values));
- let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(data)]).unwrap();
- roundtrip(batch, None);
- }
-
- test_fixed_size_binary_in_dict_inner::<UInt8Type>();
- test_fixed_size_binary_in_dict_inner::<UInt16Type>();
- test_fixed_size_binary_in_dict_inner::<UInt32Type>();
- test_fixed_size_binary_in_dict_inner::<UInt16Type>();
- test_fixed_size_binary_in_dict_inner::<Int8Type>();
- test_fixed_size_binary_in_dict_inner::<Int16Type>();
- test_fixed_size_binary_in_dict_inner::<Int32Type>();
- test_fixed_size_binary_in_dict_inner::<Int64Type>();
- }
-
- #[test]
- fn test_empty_dict() {
- let struct_fields = Fields::from(vec![Field::new(
- "dict",
- DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8)),
- false,
- )]);
-
- let schema = Schema::new(vec![Field::new_struct(
- "struct",
- struct_fields.clone(),
- true,
- )]);
- let dictionary = Arc::new(DictionaryArray::new(
- Int32Array::new_null(5),
- Arc::new(StringArray::new_null(0)),
- ));
-
- let s = StructArray::new(
- struct_fields,
- vec![dictionary],
- Some(NullBuffer::new_null(5)),
- );
-
- let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(s)]).unwrap();
- roundtrip(batch, None);
- }
#[test]
fn arrow_writer_page_size() {
let schema = Arc::new(Schema::new(vec![Field::new("col",
DataType::Utf8, false)]));
@@ -3738,140 +3114,6 @@ mod tests {
check_bloom_filter(files, "col".to_string(), optional_raw_values,
vec![""]);
}
- #[test]
- fn null_list_single_column() {
- let null_field = Field::new_list_field(DataType::Null, true);
- let list_field = Field::new("emptylist",
DataType::List(Arc::new(null_field)), true);
-
- let schema = Schema::new(vec![list_field]);
-
- // Build [[], null, [null, null]]
- let a_values = NullArray::new(2);
- let a_value_offsets = arrow::buffer::Buffer::from([0, 0, 0,
2].to_byte_slice());
- let a_list_data =
ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
- DataType::Null,
- true,
- ))))
- .len(3)
- .add_buffer(a_value_offsets)
- .null_bit_buffer(Some(Buffer::from([0b00000101])))
- .add_child_data(a_values.into_data())
- .build()
- .unwrap();
-
- let a = ListArray::from(a_list_data);
-
- assert!(a.is_valid(0));
- assert!(!a.is_valid(1));
- assert!(a.is_valid(2));
-
- assert_eq!(a.value(0).len(), 0);
- assert_eq!(a.value(2).len(), 2);
- assert_eq!(a.value(2).logical_nulls().unwrap().null_count(), 2);
-
- let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(a)]).unwrap();
- roundtrip(batch, None);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn list_single_column() {
- let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
- let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6,
10].to_byte_slice());
- let a_list_data =
ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
- DataType::Int32,
- false,
- ))))
- .len(5)
- .add_buffer(a_value_offsets)
- .null_bit_buffer(Some(Buffer::from([0b00011011])))
- .add_child_data(a_values.into_data())
- .build()
- .unwrap();
-
- assert_eq!(a_list_data.null_count(), 1);
-
- let a = ListArray::from(a_list_data);
- let values = Arc::new(a);
-
- RoundTripTest::new(values).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn large_list_single_column() {
- let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
- let a_value_offsets = arrow::buffer::Buffer::from([0i64, 1, 3, 3, 6,
10].to_byte_slice());
- let a_list_data =
ArrayData::builder(DataType::LargeList(Arc::new(Field::new(
- "large_item",
- DataType::Int32,
- true,
- ))))
- .len(5)
- .add_buffer(a_value_offsets)
- .add_child_data(a_values.into_data())
- .null_bit_buffer(Some(Buffer::from([0b00011011])))
- .build()
- .unwrap();
-
- // I think this setup is incorrect because this should pass
- assert_eq!(a_list_data.null_count(), 1);
-
- let a = LargeListArray::from(a_list_data);
- let values = Arc::new(a);
-
- RoundTripTest::new(values).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn list_nested_nulls() {
- use arrow::datatypes::Int32Type;
- let data = vec![
- Some(vec![Some(1)]),
- Some(vec![Some(2), Some(3)]),
- None,
- Some(vec![Some(4), Some(5), None]),
- Some(vec![None]),
- Some(vec![Some(6), Some(7)]),
- ];
-
- let list = ListArray::from_iter_primitive::<Int32Type, _,
_>(data.clone());
- RoundTripTest::new(Arc::new(list)).run();
-
- let list = LargeListArray::from_iter_primitive::<Int32Type, _,
_>(data);
- RoundTripTest::new(Arc::new(list)).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn list_utf8_view_selective_padding_roundtrip() {
- let item = Arc::new(Field::new_list_field(DataType::Utf8View, true));
- let mut builder =
ListBuilder::new(StringViewBuilder::new()).with_field(item);
- builder.values().append_value("a");
- builder.values().append_null();
- builder.append(true);
- // The null parent list covers selective padding dropping values below
- // the list definition level while preserving the preceding item null.
- builder.append(false);
- // The long string covers the non-inlined Utf8View buffer path.
- builder.values().append_value("large payload over 12 bytes");
- builder.append(true);
-
- RoundTripTest::new(Arc::new(builder.finish())).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn struct_single_column() {
- let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
- let struct_field_a = Arc::new(Field::new("f", DataType::Int32, false));
- let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as
ArrayRef)]);
-
- let values = Arc::new(s);
- RoundTripTest::new(values).with_nullable(false).run();
- }
-
#[test]
fn list_and_map_coerced_names() {
// Create map and list with non-Parquet naming
@@ -3969,29 +3211,6 @@ mod tests {
}
}
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_string_dictionary() {
- // define schema
- #[expect(deprecated)]
- let schema = Arc::new(Schema::new(vec![Field::new_dict(
- "dictionary",
- DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8)),
- true,
- 42,
- true,
- )]));
-
- // create some data
- let d: Int32DictionaryArray = [Some("alpha"), None, Some("beta"),
Some("alpha")]
- .iter()
- .copied()
- .collect();
-
- // build a record batch
- RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
- }
-
#[test]
fn arrow_writer_test_type_compatibility() {
fn ensure_compatible_write<T1, T2>(array1: T1, array2: T2,
expected_result: T1)
@@ -4223,144 +3442,6 @@ mod tests {
);
}
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_primitive_dictionary() {
- // define schema
- #[expect(deprecated)]
- let schema = Arc::new(Schema::new(vec![Field::new_dict(
- "dictionary",
- DataType::Dictionary(Box::new(DataType::UInt8),
Box::new(DataType::UInt32)),
- true,
- 42,
- true,
- )]));
-
- // create some data
- let mut builder = PrimitiveDictionaryBuilder::<UInt8Type,
UInt32Type>::new();
- builder.append(12345678).unwrap();
- builder.append_null();
- builder.append(22345678).unwrap();
- builder.append(12345678).unwrap();
- let d = builder.finish();
-
- RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_decimal32_dictionary() {
- let integers = vec![12345, 56789, 34567];
-
- let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2),
Some(1)]);
-
- let values = Decimal32Array::from(integers.clone())
- .with_precision_and_scale(5, 2)
- .unwrap();
-
- let array = DictionaryArray::new(keys, Arc::new(values));
- RoundTripTest::new(Arc::new(array.clone())).run();
-
- let values = Decimal32Array::from(integers)
- .with_precision_and_scale(9, 2)
- .unwrap();
-
- let array = array.with_values(Arc::new(values));
- RoundTripTest::new(Arc::new(array)).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_decimal64_dictionary() {
- let integers = vec![12345, 56789, 34567];
-
- let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2),
Some(1)]);
-
- let values = Decimal64Array::from(integers.clone())
- .with_precision_and_scale(5, 2)
- .unwrap();
-
- let array = DictionaryArray::new(keys, Arc::new(values));
- RoundTripTest::new(Arc::new(array.clone())).run();
-
- let values = Decimal64Array::from(integers)
- .with_precision_and_scale(12, 2)
- .unwrap();
-
- let array = array.with_values(Arc::new(values));
- RoundTripTest::new(Arc::new(array)).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_decimal128_dictionary() {
- let integers = vec![12345, 56789, 34567];
-
- let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2),
Some(1)]);
-
- let values = Decimal128Array::from(integers.clone())
- .with_precision_and_scale(5, 2)
- .unwrap();
-
- let array = DictionaryArray::new(keys, Arc::new(values));
- RoundTripTest::new(Arc::new(array.clone())).run();
-
- let values = Decimal128Array::from(integers)
- .with_precision_and_scale(12, 2)
- .unwrap();
-
- let array = array.with_values(Arc::new(values));
- RoundTripTest::new(Arc::new(array)).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_decimal256_dictionary() {
- let integers = vec![
- i256::from_i128(12345),
- i256::from_i128(56789),
- i256::from_i128(34567),
- ];
-
- let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2),
Some(1)]);
-
- let values = Decimal256Array::from(integers.clone())
- .with_precision_and_scale(5, 2)
- .unwrap();
-
- let array = DictionaryArray::new(keys, Arc::new(values));
- RoundTripTest::new(Arc::new(array.clone())).run();
-
- let values = Decimal256Array::from(integers)
- .with_precision_and_scale(12, 2)
- .unwrap();
-
- let array = array.with_values(Arc::new(values));
- RoundTripTest::new(Arc::new(array)).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_string_dictionary_unsigned_index() {
- // define schema
- #[expect(deprecated)]
- let schema = Arc::new(Schema::new(vec![Field::new_dict(
- "dictionary",
- DataType::Dictionary(Box::new(DataType::UInt8),
Box::new(DataType::Utf8)),
- true,
- 42,
- true,
- )]));
-
- // create some data
- let d: UInt8DictionaryArray = [Some("alpha"), None, Some("beta"),
Some("alpha")]
- .iter()
- .copied()
- .collect();
-
- RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
- }
-
#[test]
#[cfg_attr(miri, ignore)] // Takes too long
fn u32_min_max() {
diff --git a/parquet/tests/arrow_writer/roundtrip.rs
b/parquet/tests/arrow_writer/roundtrip.rs
index 135dd0df77..c9199db609 100644
--- a/parquet/tests/arrow_writer/roundtrip.rs
+++ b/parquet/tests/arrow_writer/roundtrip.rs
@@ -17,34 +17,41 @@
//! Round-trip tests for Arrow data written to Parquet.
-use super::roundtrip_helpers::{RoundTripTest, SMALL_SIZE,
required_and_optional, values_required};
+use super::roundtrip_helpers::{
+ RoundTripTest, SMALL_SIZE, required_and_optional, roundtrip,
values_required,
+};
use std::collections::HashMap;
use std::sync::Arc;
-use arrow_array::builder::FixedSizeBinaryBuilder;
+use arrow::datatypes::ToByteSlice;
+use arrow_array::builder::{
+ FixedSizeBinaryBuilder, ListBuilder, PrimitiveDictionaryBuilder,
StringViewBuilder,
+};
use arrow_array::cast::AsArray;
use arrow_array::types::{
- Date32Type, Date64Type, Decimal32Type, Decimal64Type, Decimal128Type,
Decimal256Type,
- DecimalType, Float16Type, Time32MillisecondType, Time64MicrosecondType,
+ ArrowDictionaryKeyType, Date32Type, Date64Type, Decimal32Type,
Decimal64Type, Decimal128Type,
+ Decimal256Type, DecimalType, Float16Type, Int8Type, Int16Type, Int32Type,
Int64Type,
+ Time32MillisecondType, Time64MicrosecondType, UInt8Type, UInt16Type,
UInt32Type,
};
use arrow_array::{
- Array, ArrayRef, BinaryArray, BinaryViewArray, Date32Array, Date64Array,
Decimal128Array,
- Decimal256Array, DictionaryArray, DurationMicrosecondArray,
DurationMillisecondArray,
- DurationNanosecondArray, DurationSecondArray, FixedSizeBinaryArray,
Float16Array, Float32Array,
- Float64Array, Int8Array, Int16Array, Int32Array, Int64Array,
LargeBinaryArray,
- LargeStringArray, ListArray, PrimitiveArray, RecordBatch,
RecordBatchReader, StringArray,
- StringViewArray, StructArray, Time32MillisecondArray, Time32SecondArray,
- Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray,
+ Array, ArrayRef, BinaryArray, BinaryViewArray, Date32Array, Date64Array,
Decimal32Array,
+ Decimal64Array, Decimal128Array, Decimal256Array, DictionaryArray,
DurationMicrosecondArray,
+ DurationMillisecondArray, DurationNanosecondArray, DurationSecondArray,
FixedSizeBinaryArray,
+ Float16Array, Float32Array, Float64Array, Int8Array, Int16Array,
Int32Array,
+ Int32DictionaryArray, Int64Array, LargeBinaryArray, LargeListArray,
LargeListViewArray,
+ LargeStringArray, ListArray, ListViewArray, NullArray, PrimitiveArray,
RecordBatch,
+ RecordBatchReader, StringArray, StringViewArray, StructArray,
Time32MillisecondArray,
+ Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
TimestampMicrosecondArray,
TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray,
UInt8Array,
UInt8DictionaryArray, UInt16Array, UInt32Array, UInt64Array,
};
use arrow_buffer::{ArrowNativeType, Buffer, NullBuffer, i256};
-use arrow_data::ArrayDataBuilder;
+use arrow_data::{ArrayData, ArrayDataBuilder};
use arrow_schema::{DataType as ArrowDataType, Field, Fields, Schema, TimeUnit};
use bytes::Bytes;
use half::f16;
-use num_traits::PrimInt;
+use num_traits::{FromPrimitive, PrimInt, ToPrimitive};
use parquet::arrow::ArrowWriter;
use parquet::arrow::arrow_reader::{ParquetRecordBatchReader,
ParquetRecordBatchReaderBuilder};
use parquet::basic::Type as PhysicalType;
@@ -285,6 +292,945 @@ fn string_view_single_column() {
required_and_optional::<StringViewArray, _>(raw_strs);
}
+#[test]
+fn null_list_single_column() {
+ let null_field = Field::new_list_field(ArrowDataType::Null, true);
+ let list_field = Field::new("emptylist",
ArrowDataType::List(Arc::new(null_field)), true);
+
+ let schema = Schema::new(vec![list_field]);
+
+ // Build [[], null, [null, null]]
+ let a_values = NullArray::new(2);
+ let a_value_offsets = arrow::buffer::Buffer::from([0, 0, 0,
2].to_byte_slice());
+ let a_list_data =
ArrayData::builder(ArrowDataType::List(Arc::new(Field::new_list_field(
+ ArrowDataType::Null,
+ true,
+ ))))
+ .len(3)
+ .add_buffer(a_value_offsets)
+ .null_bit_buffer(Some(Buffer::from([0b00000101])))
+ .add_child_data(a_values.into_data())
+ .build()
+ .unwrap();
+
+ let a = ListArray::from(a_list_data);
+
+ assert!(a.is_valid(0));
+ assert!(!a.is_valid(1));
+ assert!(a.is_valid(2));
+
+ assert_eq!(a.value(0).len(), 0);
+ assert_eq!(a.value(2).len(), 2);
+ assert_eq!(a.value(2).logical_nulls().unwrap().null_count(), 2);
+
+ let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(a)]).unwrap();
+ roundtrip(batch, None);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn list_single_column() {
+ let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+ let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6,
10].to_byte_slice());
+ let a_list_data =
ArrayData::builder(ArrowDataType::List(Arc::new(Field::new_list_field(
+ ArrowDataType::Int32,
+ false,
+ ))))
+ .len(5)
+ .add_buffer(a_value_offsets)
+ .null_bit_buffer(Some(Buffer::from([0b00011011])))
+ .add_child_data(a_values.into_data())
+ .build()
+ .unwrap();
+
+ assert_eq!(a_list_data.null_count(), 1);
+
+ let a = ListArray::from(a_list_data);
+ let values = Arc::new(a);
+
+ RoundTripTest::new(values).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn large_list_single_column() {
+ let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+ let a_value_offsets = arrow::buffer::Buffer::from([0i64, 1, 3, 3, 6,
10].to_byte_slice());
+ let a_list_data =
ArrayData::builder(ArrowDataType::LargeList(Arc::new(Field::new(
+ "large_item",
+ ArrowDataType::Int32,
+ true,
+ ))))
+ .len(5)
+ .add_buffer(a_value_offsets)
+ .add_child_data(a_values.into_data())
+ .null_bit_buffer(Some(Buffer::from([0b00011011])))
+ .build()
+ .unwrap();
+
+ // I think this setup is incorrect because this should pass
+ assert_eq!(a_list_data.null_count(), 1);
+
+ let a = LargeListArray::from(a_list_data);
+ let values = Arc::new(a);
+
+ RoundTripTest::new(values).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn list_nested_nulls() {
+ use arrow::datatypes::Int32Type;
+ let data = vec![
+ Some(vec![Some(1)]),
+ Some(vec![Some(2), Some(3)]),
+ None,
+ Some(vec![Some(4), Some(5), None]),
+ Some(vec![None]),
+ Some(vec![Some(6), Some(7)]),
+ ];
+
+ let list = ListArray::from_iter_primitive::<Int32Type, _, _>(data.clone());
+ RoundTripTest::new(Arc::new(list)).run();
+
+ let list = LargeListArray::from_iter_primitive::<Int32Type, _, _>(data);
+ RoundTripTest::new(Arc::new(list)).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn list_utf8_view_selective_padding_roundtrip() {
+ let item = Arc::new(Field::new_list_field(ArrowDataType::Utf8View, true));
+ let mut builder =
ListBuilder::new(StringViewBuilder::new()).with_field(item);
+ builder.values().append_value("a");
+ builder.values().append_null();
+ builder.append(true);
+ // The null parent list covers selective padding dropping values below
+ // the list definition level while preserving the preceding item null.
+ builder.append(false);
+ // The long string covers the non-inlined Utf8View buffer path.
+ builder.values().append_value("large payload over 12 bytes");
+ builder.append(true);
+
+ RoundTripTest::new(Arc::new(builder.finish())).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn struct_single_column() {
+ let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+ let struct_field_a = Arc::new(Field::new("f", ArrowDataType::Int32,
false));
+ let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as
ArrayRef)]);
+
+ let values = Arc::new(s);
+ RoundTripTest::new(values).with_nullable(false).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_list() {
+ // define schema
+ let schema = Schema::new(vec![Field::new(
+ "a",
+
ArrowDataType::List(Arc::new(Field::new_list_field(ArrowDataType::Int32,
false))),
+ true,
+ )]);
+
+ // create some data
+ let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+
+ // Construct a buffer for value offsets, for the nested array:
+ // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
+ let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6,
10].to_byte_slice());
+
+ // Construct a list array from the above two
+ let a_list_data =
ArrayData::builder(ArrowDataType::List(Arc::new(Field::new_list_field(
+ ArrowDataType::Int32,
+ false,
+ ))))
+ .len(5)
+ .add_buffer(a_value_offsets)
+ .add_child_data(a_values.into_data())
+ .null_bit_buffer(Some(Buffer::from([0b00011011])))
+ .build()
+ .unwrap();
+ let a = ListArray::from(a_list_data);
+ assert_eq!(a.null_count(), 1);
+
+ RoundTripTest::new(Arc::new(a))
+ .with_schema(Arc::new(schema))
+ .run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_list_non_null() {
+ // define schema
+ let schema = Schema::new(vec![Field::new(
+ "a",
+
ArrowDataType::List(Arc::new(Field::new_list_field(ArrowDataType::Int32,
false))),
+ false,
+ )]);
+
+ // create some data
+ let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+
+ // Construct a buffer for value offsets, for the nested array:
+ // [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
+ let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6,
10].to_byte_slice());
+
+ // Construct a list array from the above two
+ let a_list_data =
ArrayData::builder(ArrowDataType::List(Arc::new(Field::new_list_field(
+ ArrowDataType::Int32,
+ false,
+ ))))
+ .len(5)
+ .add_buffer(a_value_offsets)
+ .add_child_data(a_values.into_data())
+ .build()
+ .unwrap();
+ let a = ListArray::from(a_list_data);
+ assert_eq!(a.null_count(), 0);
+
+ RoundTripTest::new(Arc::new(a))
+ .with_schema(Arc::new(schema))
+ .run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_list_view() {
+ let list_field = Arc::new(Field::new_list_field(ArrowDataType::Int32,
false));
+ let schema = Schema::new(vec![Field::new(
+ "a",
+ ArrowDataType::ListView(list_field.clone()),
+ true,
+ )]);
+
+ // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
+ let a = ListViewArray::new(
+ list_field,
+ vec![0, 1, 0, 3, 6].into(),
+ vec![1, 2, 0, 3, 4].into(),
+ Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
+ Some(vec![true, true, false, true, true].into()),
+ );
+ assert_eq!(a.null_count(), 1);
+
+ RoundTripTest::new(Arc::new(a))
+ .with_schema(Arc::new(schema))
+ .run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_list_view_non_null() {
+ let list_field = Arc::new(Field::new_list_field(ArrowDataType::Int32,
false));
+ let schema = Schema::new(vec![Field::new(
+ "a",
+ ArrowDataType::ListView(list_field.clone()),
+ false,
+ )]);
+
+ // [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
+ let a = ListViewArray::new(
+ list_field,
+ vec![0, 1, 0, 3, 6].into(),
+ vec![1, 2, 0, 3, 4].into(),
+ Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
+ None,
+ );
+ assert_eq!(a.null_count(), 0);
+
+ RoundTripTest::new(Arc::new(a))
+ .with_schema(Arc::new(schema))
+ .run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_list_view_out_of_order() {
+ let list_field = Arc::new(Field::new_list_field(ArrowDataType::Int32,
false));
+ let schema = Schema::new(vec![Field::new(
+ "a",
+ ArrowDataType::ListView(list_field.clone()),
+ false,
+ )]);
+
+ // [[1], [2, 3], [], [7, 8, 9, 10], [4, 5, 6]] - out of order offsets
+ let a = ListViewArray::new(
+ list_field,
+ vec![0, 1, 0, 6, 3].into(),
+ vec![1, 2, 0, 4, 3].into(),
+ Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
+ None,
+ );
+ assert_eq!(a.null_count(), 0);
+
+ RoundTripTest::new(Arc::new(a))
+ .with_schema(Arc::new(schema))
+ .run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_large_list_view() {
+ let list_field = Arc::new(Field::new_list_field(ArrowDataType::Int32,
false));
+ let schema = Schema::new(vec![Field::new(
+ "a",
+ ArrowDataType::LargeListView(list_field.clone()),
+ true,
+ )]);
+
+ // [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
+ let a = LargeListViewArray::new(
+ list_field,
+ vec![0i64, 1, 0, 3, 6].into(),
+ vec![1i64, 2, 0, 3, 4].into(),
+ Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
+ Some(vec![true, true, false, true, true].into()),
+ );
+ assert_eq!(a.null_count(), 1);
+
+ RoundTripTest::new(Arc::new(a))
+ .with_schema(Arc::new(schema))
+ .run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_list_view_with_struct() {
+ // Test ListView containing Struct: ListView<Struct<Int32, Utf8>>
+ let struct_fields = Fields::from(vec![
+ Field::new("id", ArrowDataType::Int32, false),
+ Field::new("name", ArrowDataType::Utf8, false),
+ ]);
+ let struct_type = ArrowDataType::Struct(struct_fields.clone());
+ let list_field = Arc::new(Field::new("item", struct_type.clone(), false));
+
+ let schema = Schema::new(vec![Field::new(
+ "a",
+ ArrowDataType::ListView(list_field.clone()),
+ true,
+ )]);
+
+ // Create struct values
+ let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
+ let name_array = StringArray::from(vec!["a", "b", "c", "d", "e"]);
+ let struct_array = StructArray::new(
+ struct_fields,
+ vec![Arc::new(id_array), Arc::new(name_array)],
+ None,
+ );
+
+ // Create ListView: [{1, "a"}, {2, "b"}], null, [{3, "c"}, {4, "d"}, {5,
"e"}]
+ let list_view = ListViewArray::new(
+ list_field,
+ vec![0, 2, 2].into(), // offsets
+ vec![2, 0, 3].into(), // sizes
+ Arc::new(struct_array),
+ Some(vec![true, false, true].into()),
+ );
+ assert_eq!(list_view.null_count(), 1);
+
+ RoundTripTest::new(Arc::new(list_view))
+ .with_schema(Arc::new(schema))
+ .run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_complex() {
+ // define schema
+ let struct_field_d = Arc::new(Field::new("d", ArrowDataType::Float64,
true));
+ let struct_field_f = Arc::new(Field::new("f", ArrowDataType::Float32,
true));
+ let struct_field_g = Arc::new(Field::new_list(
+ "g",
+ Field::new_list_field(ArrowDataType::Int16, true),
+ false,
+ ));
+ let struct_field_h = Arc::new(Field::new_list(
+ "h",
+ Field::new_list_field(ArrowDataType::Int16, false),
+ true,
+ ));
+ let struct_field_e = Arc::new(Field::new_struct(
+ "e",
+ vec![
+ struct_field_f.clone(),
+ struct_field_g.clone(),
+ struct_field_h.clone(),
+ ],
+ false,
+ ));
+ let schema = Schema::new(vec![
+ Field::new("a", ArrowDataType::Int32, false),
+ Field::new("b", ArrowDataType::Int32, true),
+ Field::new_struct(
+ "c",
+ vec![struct_field_d.clone(), struct_field_e.clone()],
+ false,
+ ),
+ ]);
+
+ // create some data
+ let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
+ let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
+ let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
+ let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None,
Some(5.25)]);
+
+ let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+
+ // Construct a buffer for value offsets, for the nested array:
+ // [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
+ let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6,
10].to_byte_slice());
+
+ // Construct a list array from the above two
+ let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
+ .len(5)
+ .add_buffer(g_value_offsets.clone())
+ .add_child_data(g_value.to_data())
+ .build()
+ .unwrap();
+ let g = ListArray::from(g_list_data);
+ // The difference between g and h is that h has a null bitmap
+ let h_list_data = ArrayData::builder(struct_field_h.data_type().clone())
+ .len(5)
+ .add_buffer(g_value_offsets)
+ .add_child_data(g_value.to_data())
+ .null_bit_buffer(Some(Buffer::from([0b00011011])))
+ .build()
+ .unwrap();
+ let h = ListArray::from(h_list_data);
+
+ let e = StructArray::from(vec![
+ (struct_field_f, Arc::new(f) as ArrayRef),
+ (struct_field_g, Arc::new(g) as ArrayRef),
+ (struct_field_h, Arc::new(h) as ArrayRef),
+ ]);
+
+ let c = StructArray::from(vec![
+ (struct_field_d, Arc::new(d) as ArrayRef),
+ (struct_field_e, Arc::new(e) as ArrayRef),
+ ]);
+
+ // build a record batch
+ let batch = RecordBatch::try_new(
+ Arc::new(schema),
+ vec![Arc::new(a), Arc::new(b), Arc::new(c)],
+ )
+ .unwrap();
+
+ roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
+ roundtrip(batch, Some(SMALL_SIZE / 3));
+}
+
+#[test]
+fn arrow_writer_complex_mixed() {
+ // This test was added while investigating
https://github.com/apache/arrow-rs/issues/244.
+ // It was subsequently fixed while investigating
https://github.com/apache/arrow-rs/issues/245.
+
+ // define schema
+ let offset_field = Arc::new(Field::new("offset", ArrowDataType::Int32,
false));
+ let partition_field = Arc::new(Field::new("partition",
ArrowDataType::Int64, true));
+ let topic_field = Arc::new(Field::new("topic", ArrowDataType::Utf8, true));
+ let schema = Schema::new(vec![Field::new(
+ "some_nested_object",
+ ArrowDataType::Struct(Fields::from(vec![
+ offset_field.clone(),
+ partition_field.clone(),
+ topic_field.clone(),
+ ])),
+ false,
+ )]);
+
+ // create some data
+ let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
+ let partition = Int64Array::from(vec![Some(1), None, None, Some(4),
Some(5)]);
+ let topic = StringArray::from(vec![Some("A"), None, Some("A"), Some(""),
None]);
+
+ let some_nested_object = StructArray::from(vec![
+ (offset_field, Arc::new(offset) as ArrayRef),
+ (partition_field, Arc::new(partition) as ArrayRef),
+ (topic_field, Arc::new(topic) as ArrayRef),
+ ]);
+
+ // build a record batch
+ let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(some_nested_object)]).unwrap();
+
+ roundtrip(batch, Some(SMALL_SIZE / 2));
+}
+
+#[test]
+fn arrow_writer_map() {
+ // Note: we are using the JSON Arrow reader for brevity
+ let json_content = r#"
+ {"stocks":{"long": "$AAA", "short": "$BBB"}}
+ {"stocks":{"long": null, "long": "$CCC", "short": null}}
+ {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}}
+ "#;
+ let entries_struct_type = ArrowDataType::Struct(Fields::from(vec![
+ Field::new(
+ Field::MAP_KEY_FIELD_DEFAULT_NAME,
+ ArrowDataType::Utf8,
+ false,
+ ),
+ Field::new(
+ Field::MAP_VALUE_FIELD_DEFAULT_NAME,
+ ArrowDataType::Utf8,
+ true,
+ ),
+ ]));
+ let stocks_field = Field::new(
+ "stocks",
+ ArrowDataType::Map(
+ Arc::new(Field::new(
+ Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
+ entries_struct_type,
+ false,
+ )),
+ false,
+ ),
+ true,
+ );
+ let schema = Arc::new(Schema::new(vec![stocks_field]));
+ let builder = arrow::json::ReaderBuilder::new(schema).with_batch_size(64);
+ let mut reader =
builder.build(std::io::Cursor::new(json_content)).unwrap();
+
+ let batch = reader.next().unwrap().unwrap();
+ roundtrip(batch, None);
+}
+
+#[test]
+fn arrow_writer_2_level_struct() {
+ // tests writing <struct<struct<primitive>>
+ let field_c = Field::new("c", ArrowDataType::Int32, true);
+ let field_b = Field::new("b", ArrowDataType::Struct(vec![field_c].into()),
true);
+ let type_a = ArrowDataType::Struct(vec![field_b.clone()].into());
+ let field_a = Field::new("a", type_a, true);
+ let schema = Schema::new(vec![field_a.clone()]);
+
+ // create data
+ let c = Int32Array::from(vec![Some(1), None, Some(3), None, None,
Some(6)]);
+ let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
+ .len(6)
+ .null_bit_buffer(Some(Buffer::from([0b00100111])))
+ .add_child_data(c.into_data())
+ .build()
+ .unwrap();
+ let b = StructArray::from(b_data);
+ let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
+ .len(6)
+ .null_bit_buffer(Some(Buffer::from([0b00101111])))
+ .add_child_data(b.into_data())
+ .build()
+ .unwrap();
+ let a = StructArray::from(a_data);
+
+ assert_eq!(a.null_count(), 1);
+ assert_eq!(a.column(0).null_count(), 2);
+
+ // build a racord batch
+ let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(a)]).unwrap();
+
+ roundtrip(batch, Some(SMALL_SIZE / 2));
+}
+
+#[test]
+fn arrow_writer_2_level_struct_non_null() {
+ // tests writing <struct<struct<primitive>>
+ let field_c = Field::new("c", ArrowDataType::Int32, false);
+ let type_b = ArrowDataType::Struct(vec![field_c].into());
+ let field_b = Field::new("b", type_b.clone(), false);
+ let type_a = ArrowDataType::Struct(vec![field_b].into());
+ let field_a = Field::new("a", type_a.clone(), false);
+ let schema = Schema::new(vec![field_a]);
+
+ // create data
+ let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
+ let b_data = ArrayDataBuilder::new(type_b)
+ .len(6)
+ .add_child_data(c.into_data())
+ .build()
+ .unwrap();
+ let b = StructArray::from(b_data);
+ let a_data = ArrayDataBuilder::new(type_a)
+ .len(6)
+ .add_child_data(b.into_data())
+ .build()
+ .unwrap();
+ let a = StructArray::from(a_data);
+
+ assert_eq!(a.null_count(), 0);
+ assert_eq!(a.column(0).null_count(), 0);
+
+ // build a racord batch
+ let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(a)]).unwrap();
+
+ roundtrip(batch, Some(SMALL_SIZE / 2));
+}
+
+#[test]
+fn arrow_writer_2_level_struct_mixed_null() {
+ // tests writing <struct<struct<primitive>>
+ let field_c = Field::new("c", ArrowDataType::Int32, false);
+ let type_b = ArrowDataType::Struct(vec![field_c].into());
+ let field_b = Field::new("b", type_b.clone(), true);
+ let type_a = ArrowDataType::Struct(vec![field_b].into());
+ let field_a = Field::new("a", type_a.clone(), false);
+ let schema = Schema::new(vec![field_a]);
+
+ // create data
+ let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
+ let b_data = ArrayDataBuilder::new(type_b)
+ .len(6)
+ .null_bit_buffer(Some(Buffer::from([0b00100111])))
+ .add_child_data(c.into_data())
+ .build()
+ .unwrap();
+ let b = StructArray::from(b_data);
+ // a intentionally has no null buffer, to test that this is handled
correctly
+ let a_data = ArrayDataBuilder::new(type_a)
+ .len(6)
+ .add_child_data(b.into_data())
+ .build()
+ .unwrap();
+ let a = StructArray::from(a_data);
+
+ assert_eq!(a.null_count(), 0);
+ assert_eq!(a.column(0).null_count(), 2);
+
+ // build a racord batch
+ let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(a)]).unwrap();
+
+ roundtrip(batch, Some(SMALL_SIZE / 2));
+}
+
+#[test]
+fn arrow_writer_2_level_struct_mixed_null_2() {
+ // tests writing <struct<struct<primitive>>, where the primitive columns
are non-null.
+ let field_c = Field::new("c", ArrowDataType::Int32, false);
+ let field_d = Field::new("d", ArrowDataType::FixedSizeBinary(4), false);
+ let field_e = Field::new(
+ "e",
+ ArrowDataType::Dictionary(
+ Box::new(ArrowDataType::Int32),
+ Box::new(ArrowDataType::Utf8),
+ ),
+ false,
+ );
+
+ let field_b = Field::new(
+ "b",
+ ArrowDataType::Struct(vec![field_c, field_d, field_e].into()),
+ false,
+ );
+ let type_a = ArrowDataType::Struct(vec![field_b.clone()].into());
+ let field_a = Field::new("a", type_a, true);
+ let schema = Schema::new(vec![field_a.clone()]);
+
+ // create data
+ let c = Int32Array::from_iter_values(0..6);
+ let d = FixedSizeBinaryArray::try_from_iter(
+ ["aaaa", "bbbb", "cccc", "dddd", "eeee", "ffff"].into_iter(),
+ )
+ .expect("four byte values");
+ let e = Int32DictionaryArray::from_iter(["one", "two", "three", "four",
"five", "one"]);
+ let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
+ .len(6)
+ .add_child_data(c.into_data())
+ .add_child_data(d.into_data())
+ .add_child_data(e.into_data())
+ .build()
+ .unwrap();
+ let b = StructArray::from(b_data);
+ let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
+ .len(6)
+ .null_bit_buffer(Some(Buffer::from([0b00100101])))
+ .add_child_data(b.into_data())
+ .build()
+ .unwrap();
+ let a = StructArray::from(a_data);
+
+ assert_eq!(a.null_count(), 3);
+ assert_eq!(a.column(0).null_count(), 0);
+
+ // build a record batch
+ let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(a)]).unwrap();
+
+ roundtrip(batch, Some(SMALL_SIZE / 2));
+}
+
+/// Test round-trip of Dictionary<UInt32, Utf8View> and
+/// Dictionary<UInt32, BinaryView> typed columns.
+#[test]
+fn arrow_writer_string_view_dictionary() {
+ let raw_string_values = vec!["a", "b", "large payload over 12 bytes"];
+ let raw_binary_values = vec![
+ b"a".to_vec(),
+ b"b".to_vec(),
+ b"large payload over 12 bytes".to_vec(),
+ ];
+
+ let keys = UInt32Array::from(vec![Some(0), None, Some(2), Some(1), None]);
+
+ let string_view_values =
Arc::new(StringViewArray::from(raw_string_values));
+ let string_dict: ArrayRef =
+ Arc::new(DictionaryArray::<UInt32Type>::try_new(keys.clone(),
string_view_values).unwrap());
+
+ let binary_view_values =
Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
+ let binary_dict: ArrayRef =
+ Arc::new(DictionaryArray::<UInt32Type>::try_new(keys,
binary_view_values).unwrap());
+
+ RoundTripTest::new(string_dict).run();
+ RoundTripTest::new(binary_dict).run();
+}
+
+#[test]
+fn test_fixed_size_binary_in_dict() {
+ fn test_fixed_size_binary_in_dict_inner<K>()
+ where
+ K: ArrowDictionaryKeyType,
+ K::Native: FromPrimitive + ToPrimitive + TryFrom<u8>,
+ <<K as arrow_array::ArrowPrimitiveType>::Native as
TryFrom<u8>>::Error: std::fmt::Debug,
+ {
+ let field = Field::new(
+ "a",
+ ArrowDataType::Dictionary(
+ Box::new(K::DATA_TYPE),
+ Box::new(ArrowDataType::FixedSizeBinary(4)),
+ ),
+ false,
+ );
+ let schema = Schema::new(vec![field]);
+
+ let keys: Vec<K::Native> = vec![
+ K::Native::try_from(0u8).unwrap(),
+ K::Native::try_from(0u8).unwrap(),
+ K::Native::try_from(1u8).unwrap(),
+ ];
+ let keys = PrimitiveArray::<K>::from_iter_values(keys);
+ let values = FixedSizeBinaryArray::try_from_iter(
+ vec![vec![0, 0, 0, 0], vec![1, 1, 1, 1]].into_iter(),
+ )
+ .unwrap();
+
+ let data = DictionaryArray::<K>::new(keys, Arc::new(values));
+ let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(data)]).unwrap();
+ roundtrip(batch, None);
+ }
+
+ test_fixed_size_binary_in_dict_inner::<UInt8Type>();
+ test_fixed_size_binary_in_dict_inner::<UInt16Type>();
+ test_fixed_size_binary_in_dict_inner::<UInt32Type>();
+ test_fixed_size_binary_in_dict_inner::<UInt16Type>();
+ test_fixed_size_binary_in_dict_inner::<Int8Type>();
+ test_fixed_size_binary_in_dict_inner::<Int16Type>();
+ test_fixed_size_binary_in_dict_inner::<Int32Type>();
+ test_fixed_size_binary_in_dict_inner::<Int64Type>();
+}
+
+#[test]
+fn test_empty_dict() {
+ let struct_fields = Fields::from(vec![Field::new(
+ "dict",
+ ArrowDataType::Dictionary(
+ Box::new(ArrowDataType::Int32),
+ Box::new(ArrowDataType::Utf8),
+ ),
+ false,
+ )]);
+
+ let schema = Schema::new(vec![Field::new_struct(
+ "struct",
+ struct_fields.clone(),
+ true,
+ )]);
+ let dictionary = Arc::new(DictionaryArray::new(
+ Int32Array::new_null(5),
+ Arc::new(StringArray::new_null(0)),
+ ));
+
+ let s = StructArray::new(
+ struct_fields,
+ vec![dictionary],
+ Some(NullBuffer::new_null(5)),
+ );
+
+ let batch = RecordBatch::try_new(Arc::new(schema),
vec![Arc::new(s)]).unwrap();
+ roundtrip(batch, None);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_string_dictionary() {
+ // define schema
+ #[expect(deprecated)]
+ let schema = Arc::new(Schema::new(vec![Field::new_dict(
+ "dictionary",
+ ArrowDataType::Dictionary(
+ Box::new(ArrowDataType::Int32),
+ Box::new(ArrowDataType::Utf8),
+ ),
+ true,
+ 42,
+ true,
+ )]));
+
+ // create some data
+ let d: Int32DictionaryArray = [Some("alpha"), None, Some("beta"),
Some("alpha")]
+ .iter()
+ .copied()
+ .collect();
+
+ // build a record batch
+ RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_primitive_dictionary() {
+ // define schema
+ #[expect(deprecated)]
+ let schema = Arc::new(Schema::new(vec![Field::new_dict(
+ "dictionary",
+ ArrowDataType::Dictionary(
+ Box::new(ArrowDataType::UInt8),
+ Box::new(ArrowDataType::UInt32),
+ ),
+ true,
+ 42,
+ true,
+ )]));
+
+ // create some data
+ let mut builder = PrimitiveDictionaryBuilder::<UInt8Type,
UInt32Type>::new();
+ builder.append(12345678).unwrap();
+ builder.append_null();
+ builder.append(22345678).unwrap();
+ builder.append(12345678).unwrap();
+ let d = builder.finish();
+
+ RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_decimal32_dictionary() {
+ let integers = vec![12345, 56789, 34567];
+
+ let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2),
Some(1)]);
+
+ let values = Decimal32Array::from(integers.clone())
+ .with_precision_and_scale(5, 2)
+ .unwrap();
+
+ let array = DictionaryArray::new(keys, Arc::new(values));
+ RoundTripTest::new(Arc::new(array.clone())).run();
+
+ let values = Decimal32Array::from(integers)
+ .with_precision_and_scale(9, 2)
+ .unwrap();
+
+ let array = array.with_values(Arc::new(values));
+ RoundTripTest::new(Arc::new(array)).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_decimal64_dictionary() {
+ let integers = vec![12345, 56789, 34567];
+
+ let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2),
Some(1)]);
+
+ let values = Decimal64Array::from(integers.clone())
+ .with_precision_and_scale(5, 2)
+ .unwrap();
+
+ let array = DictionaryArray::new(keys, Arc::new(values));
+ RoundTripTest::new(Arc::new(array.clone())).run();
+
+ let values = Decimal64Array::from(integers)
+ .with_precision_and_scale(12, 2)
+ .unwrap();
+
+ let array = array.with_values(Arc::new(values));
+ RoundTripTest::new(Arc::new(array)).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_decimal128_dictionary() {
+ let integers = vec![12345, 56789, 34567];
+
+ let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2),
Some(1)]);
+
+ let values = Decimal128Array::from(integers.clone())
+ .with_precision_and_scale(5, 2)
+ .unwrap();
+
+ let array = DictionaryArray::new(keys, Arc::new(values));
+ RoundTripTest::new(Arc::new(array.clone())).run();
+
+ let values = Decimal128Array::from(integers)
+ .with_precision_and_scale(12, 2)
+ .unwrap();
+
+ let array = array.with_values(Arc::new(values));
+ RoundTripTest::new(Arc::new(array)).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_decimal256_dictionary() {
+ let integers = vec![
+ i256::from_i128(12345),
+ i256::from_i128(56789),
+ i256::from_i128(34567),
+ ];
+
+ let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2),
Some(1)]);
+
+ let values = Decimal256Array::from(integers.clone())
+ .with_precision_and_scale(5, 2)
+ .unwrap();
+
+ let array = DictionaryArray::new(keys, Arc::new(values));
+ RoundTripTest::new(Arc::new(array.clone())).run();
+
+ let values = Decimal256Array::from(integers)
+ .with_precision_and_scale(12, 2)
+ .unwrap();
+
+ let array = array.with_values(Arc::new(values));
+ RoundTripTest::new(Arc::new(array)).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_string_dictionary_unsigned_index() {
+ // define schema
+ #[expect(deprecated)]
+ let schema = Arc::new(Schema::new(vec![Field::new_dict(
+ "dictionary",
+ ArrowDataType::Dictionary(
+ Box::new(ArrowDataType::UInt8),
+ Box::new(ArrowDataType::Utf8),
+ ),
+ true,
+ 42,
+ true,
+ )]));
+
+ // create some data
+ let d: UInt8DictionaryArray = [Some("alpha"), None, Some("beta"),
Some("alpha")]
+ .iter()
+ .copied()
+ .collect();
+
+ RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
+}
+
#[test]
fn test_unsigned_roundtrip() {
let schema = Arc::new(Schema::new(vec