diff --git a/arrow/src/array/array_union.rs b/arrow/src/array/array_union.rs index 09cc1ab59fe..58444b9a8fd 100644 --- a/arrow/src/array/array_union.rs +++ b/arrow/src/array/array_union.rs @@ -387,7 +387,7 @@ mod tests { #[test] fn test_dense_i32() { - let mut builder = UnionBuilder::new_dense(7); + let mut builder = UnionBuilder::new_dense(); builder.append::("a", 1).unwrap(); builder.append::("b", 2).unwrap(); builder.append::("c", 3).unwrap(); @@ -447,7 +447,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_dense_i32_large() { - let mut builder = UnionBuilder::new_dense(1024); + let mut builder = UnionBuilder::new_dense(); let expected_type_ids = vec![0_i8; 1024]; let expected_value_offsets: Vec<_> = (0..1024).collect(); @@ -489,7 +489,7 @@ mod tests { #[test] fn test_dense_mixed() { - let mut builder = UnionBuilder::new_dense(7); + let mut builder = UnionBuilder::new_dense(); builder.append::("a", 1).unwrap(); builder.append::("c", 3).unwrap(); builder.append::("a", 4).unwrap(); @@ -539,7 +539,7 @@ mod tests { #[test] fn test_dense_mixed_with_nulls() { - let mut builder = UnionBuilder::new_dense(7); + let mut builder = UnionBuilder::new_dense(); builder.append::("a", 1).unwrap(); builder.append::("c", 3).unwrap(); builder.append::("a", 10).unwrap(); @@ -587,7 +587,7 @@ mod tests { #[test] fn test_dense_mixed_with_nulls_and_offset() { - let mut builder = UnionBuilder::new_dense(7); + let mut builder = UnionBuilder::new_dense(); builder.append::("a", 1).unwrap(); builder.append::("c", 3).unwrap(); builder.append::("a", 10).unwrap(); @@ -714,7 +714,7 @@ mod tests { #[test] fn test_sparse_i32() { - let mut builder = UnionBuilder::new_sparse(7); + let mut builder = UnionBuilder::new_sparse(); builder.append::("a", 1).unwrap(); builder.append::("b", 2).unwrap(); builder.append::("c", 3).unwrap(); @@ -766,7 +766,7 @@ mod tests { #[test] fn test_sparse_mixed() { - let mut builder = UnionBuilder::new_sparse(5); + let mut builder = UnionBuilder::new_sparse(); builder.append::("a", 1).unwrap(); builder.append::("c", 3.0).unwrap(); builder.append::("a", 4).unwrap(); @@ -829,7 +829,7 @@ mod tests { #[test] fn test_sparse_mixed_with_nulls() { - let mut builder = UnionBuilder::new_sparse(5); + let mut builder = UnionBuilder::new_sparse(); builder.append::("a", 1).unwrap(); builder.append_null::("a").unwrap(); builder.append::("c", 3.0).unwrap(); @@ -882,7 +882,7 @@ mod tests { #[test] fn test_sparse_mixed_with_nulls_and_offset() { - let mut builder = UnionBuilder::new_sparse(5); + let mut builder = UnionBuilder::new_sparse(); builder.append::("a", 1).unwrap(); builder.append_null::("a").unwrap(); builder.append::("c", 3.0).unwrap(); @@ -929,7 +929,7 @@ mod tests { #[test] fn test_union_array_validaty() { - let mut builder = UnionBuilder::new_sparse(5); + let mut builder = UnionBuilder::new_sparse(); builder.append::("a", 1).unwrap(); builder.append_null::("a").unwrap(); builder.append::("c", 3.0).unwrap(); @@ -939,7 +939,7 @@ mod tests { test_union_validity(&union); - let mut builder = UnionBuilder::new_dense(5); + let mut builder = UnionBuilder::new_dense(); builder.append::("a", 1).unwrap(); builder.append_null::("a").unwrap(); builder.append::("c", 3.0).unwrap(); @@ -952,7 +952,7 @@ mod tests { #[test] fn test_type_check() { - let mut builder = UnionBuilder::new_sparse(2); + let mut builder = UnionBuilder::new_sparse(); builder.append::("a", 1.0).unwrap(); let err = builder.append::("a", 1).unwrap_err().to_string(); assert!(err.contains("Attempt to write col \"a\" with type Int32 doesn't match existing type Float32"), "{}", err); @@ -1009,14 +1009,14 @@ mod tests { } // Sparse Union - let builder = UnionBuilder::new_sparse(5); + let builder = UnionBuilder::new_sparse(); let record_batch = create_batch(create_union(builder)); // [null, 3.0, null] let record_batch_slice = record_batch.slice(1, 3); test_slice_union(record_batch_slice); // Dense Union - let builder = UnionBuilder::new_dense(5); + let builder = UnionBuilder::new_dense(); let record_batch = create_batch(create_union(builder)); // [null, 3.0, null] let record_batch_slice = record_batch.slice(1, 3); diff --git a/arrow/src/array/builder/union_builder.rs b/arrow/src/array/builder/union_builder.rs index 8ef5028afc8..f263fecf6f2 100644 --- a/arrow/src/array/builder/union_builder.rs +++ b/arrow/src/array/builder/union_builder.rs @@ -111,7 +111,7 @@ impl FieldData { /// use arrow::array::UnionBuilder; /// use arrow::datatypes::{Float64Type, Int32Type}; /// -/// let mut builder = UnionBuilder::new_dense(3); +/// let mut builder = UnionBuilder::new_dense(); /// builder.append::("a", 1).unwrap(); /// builder.append::("b", 3.0).unwrap(); /// builder.append::("a", 4).unwrap(); @@ -131,7 +131,7 @@ impl FieldData { /// use arrow::array::UnionBuilder; /// use arrow::datatypes::{Float64Type, Int32Type}; /// -/// let mut builder = UnionBuilder::new_sparse(3); +/// let mut builder = UnionBuilder::new_sparse(); /// builder.append::("a", 1).unwrap(); /// builder.append::("b", 3.0).unwrap(); /// builder.append::("a", 4).unwrap(); @@ -159,7 +159,17 @@ pub struct UnionBuilder { impl UnionBuilder { /// Creates a new dense array builder. - pub fn new_dense(capacity: usize) -> Self { + pub fn new_dense() -> Self { + Self::with_capacity_dense(1024) + } + + /// Creates a new sparse array builder. + pub fn new_sparse() -> Self { + Self::with_capacity_sparse(1024) + } + + /// Creates a new dense array builder with capacity. + pub fn with_capacity_dense(capacity: usize) -> Self { Self { len: 0, fields: HashMap::default(), @@ -168,8 +178,8 @@ impl UnionBuilder { } } - /// Creates a new sparse array builder. - pub fn new_sparse(capacity: usize) -> Self { + /// Creates a new sparse array builder with capacity. + pub fn with_capacity_sparse(capacity: usize) -> Self { Self { len: 0, fields: HashMap::default(), diff --git a/arrow/src/array/equal/mod.rs b/arrow/src/array/equal/mod.rs index 975fbd2b360..debb8e7240c 100644 --- a/arrow/src/array/equal/mod.rs +++ b/arrow/src/array/equal/mod.rs @@ -1370,7 +1370,7 @@ mod tests { #[test] fn test_union_equal_dense() { - let mut builder = UnionBuilder::new_dense(7); + let mut builder = UnionBuilder::new_dense(); builder.append::("a", 1).unwrap(); builder.append::("b", 2).unwrap(); builder.append::("c", 3).unwrap(); @@ -1380,7 +1380,7 @@ mod tests { builder.append::("b", 7).unwrap(); let union1 = builder.build().unwrap(); - builder = UnionBuilder::new_dense(7); + builder = UnionBuilder::new_dense(); builder.append::("a", 1).unwrap(); builder.append::("b", 2).unwrap(); builder.append::("c", 3).unwrap(); @@ -1390,7 +1390,7 @@ mod tests { builder.append::("b", 7).unwrap(); let union2 = builder.build().unwrap(); - builder = UnionBuilder::new_dense(7); + builder = UnionBuilder::new_dense(); builder.append::("a", 1).unwrap(); builder.append::("b", 2).unwrap(); builder.append::("c", 3).unwrap(); @@ -1400,7 +1400,7 @@ mod tests { builder.append::("b", 7).unwrap(); let union3 = builder.build().unwrap(); - builder = UnionBuilder::new_dense(7); + builder = UnionBuilder::new_dense(); builder.append::("a", 1).unwrap(); builder.append::("b", 2).unwrap(); builder.append::("c", 3).unwrap(); @@ -1417,7 +1417,7 @@ mod tests { #[test] fn test_union_equal_sparse() { - let mut builder = UnionBuilder::new_sparse(7); + let mut builder = UnionBuilder::new_sparse(); builder.append::("a", 1).unwrap(); builder.append::("b", 2).unwrap(); builder.append::("c", 3).unwrap(); @@ -1427,7 +1427,7 @@ mod tests { builder.append::("b", 7).unwrap(); let union1 = builder.build().unwrap(); - builder = UnionBuilder::new_sparse(7); + builder = UnionBuilder::new_sparse(); builder.append::("a", 1).unwrap(); builder.append::("b", 2).unwrap(); builder.append::("c", 3).unwrap(); @@ -1437,7 +1437,7 @@ mod tests { builder.append::("b", 7).unwrap(); let union2 = builder.build().unwrap(); - builder = UnionBuilder::new_sparse(7); + builder = UnionBuilder::new_sparse(); builder.append::("a", 1).unwrap(); builder.append::("b", 2).unwrap(); builder.append::("c", 3).unwrap(); @@ -1447,7 +1447,7 @@ mod tests { builder.append::("b", 7).unwrap(); let union3 = builder.build().unwrap(); - builder = UnionBuilder::new_sparse(7); + builder = UnionBuilder::new_sparse(); builder.append::("a", 1).unwrap(); builder.append::("b", 2).unwrap(); builder.append::("c", 3).unwrap(); diff --git a/arrow/src/compute/kernels/cast.rs b/arrow/src/compute/kernels/cast.rs index 4f59c00fcb5..2d9e0b5ded7 100644 --- a/arrow/src/compute/kernels/cast.rs +++ b/arrow/src/compute/kernels/cast.rs @@ -5367,7 +5367,7 @@ mod tests { } fn make_union_array() -> UnionArray { - let mut builder = UnionBuilder::new_dense(7); + let mut builder = UnionBuilder::with_capacity_dense(7); builder.append::("a", 1).unwrap(); builder.append::("b", 2).unwrap(); builder.build().unwrap() diff --git a/arrow/src/compute/kernels/filter.rs b/arrow/src/compute/kernels/filter.rs index 95a1bfa4da4..0bffee62cd9 100644 --- a/arrow/src/compute/kernels/filter.rs +++ b/arrow/src/compute/kernels/filter.rs @@ -1553,7 +1553,7 @@ mod tests { let c = filter(&array, &filter_array).unwrap(); let filtered = c.as_any().downcast_ref::().unwrap(); - let mut builder = UnionBuilder::new_dense(1); + let mut builder = UnionBuilder::new_dense(); builder.append::("A", 1).unwrap(); let expected_array = builder.build().unwrap(); @@ -1563,7 +1563,7 @@ mod tests { let c = filter(&array, &filter_array).unwrap(); let filtered = c.as_any().downcast_ref::().unwrap(); - let mut builder = UnionBuilder::new_dense(2); + let mut builder = UnionBuilder::new_dense(); builder.append::("A", 1).unwrap(); builder.append::("A", 34).unwrap(); let expected_array = builder.build().unwrap(); @@ -1574,7 +1574,7 @@ mod tests { let c = filter(&array, &filter_array).unwrap(); let filtered = c.as_any().downcast_ref::().unwrap(); - let mut builder = UnionBuilder::new_dense(2); + let mut builder = UnionBuilder::new_dense(); builder.append::("A", 1).unwrap(); builder.append::("B", 3.2).unwrap(); let expected_array = builder.build().unwrap(); @@ -1584,7 +1584,7 @@ mod tests { #[test] fn test_filter_union_array_dense() { - let mut builder = UnionBuilder::new_dense(3); + let mut builder = UnionBuilder::new_dense(); builder.append::("A", 1).unwrap(); builder.append::("B", 3.2).unwrap(); builder.append::("A", 34).unwrap(); @@ -1595,7 +1595,7 @@ mod tests { #[test] fn test_filter_run_union_array_dense() { - let mut builder = UnionBuilder::new_dense(3); + let mut builder = UnionBuilder::new_dense(); builder.append::("A", 1).unwrap(); builder.append::("A", 3).unwrap(); builder.append::("A", 34).unwrap(); @@ -1605,7 +1605,7 @@ mod tests { let c = filter(&array, &filter_array).unwrap(); let filtered = c.as_any().downcast_ref::().unwrap(); - let mut builder = UnionBuilder::new_dense(3); + let mut builder = UnionBuilder::new_dense(); builder.append::("A", 1).unwrap(); builder.append::("A", 3).unwrap(); let expected = builder.build().unwrap(); @@ -1615,7 +1615,7 @@ mod tests { #[test] fn test_filter_union_array_dense_with_nulls() { - let mut builder = UnionBuilder::new_dense(4); + let mut builder = UnionBuilder::new_dense(); builder.append::("A", 1).unwrap(); builder.append::("B", 3.2).unwrap(); builder.append_null::("B").unwrap(); @@ -1626,7 +1626,7 @@ mod tests { let c = filter(&array, &filter_array).unwrap(); let filtered = c.as_any().downcast_ref::().unwrap(); - let mut builder = UnionBuilder::new_dense(2); + let mut builder = UnionBuilder::new_dense(); builder.append::("A", 1).unwrap(); builder.append::("B", 3.2).unwrap(); let expected_array = builder.build().unwrap(); @@ -1637,7 +1637,7 @@ mod tests { let c = filter(&array, &filter_array).unwrap(); let filtered = c.as_any().downcast_ref::().unwrap(); - let mut builder = UnionBuilder::new_dense(2); + let mut builder = UnionBuilder::new_dense(); builder.append::("A", 1).unwrap(); builder.append_null::("B").unwrap(); let expected_array = builder.build().unwrap(); @@ -1647,7 +1647,7 @@ mod tests { #[test] fn test_filter_union_array_sparse() { - let mut builder = UnionBuilder::new_sparse(3); + let mut builder = UnionBuilder::new_sparse(); builder.append::("A", 1).unwrap(); builder.append::("B", 3.2).unwrap(); builder.append::("A", 34).unwrap(); @@ -1658,7 +1658,7 @@ mod tests { #[test] fn test_filter_union_array_sparse_with_nulls() { - let mut builder = UnionBuilder::new_sparse(4); + let mut builder = UnionBuilder::new_sparse(); builder.append::("A", 1).unwrap(); builder.append::("B", 3.2).unwrap(); builder.append_null::("B").unwrap(); @@ -1669,7 +1669,7 @@ mod tests { let c = filter(&array, &filter_array).unwrap(); let filtered = c.as_any().downcast_ref::().unwrap(); - let mut builder = UnionBuilder::new_sparse(2); + let mut builder = UnionBuilder::new_sparse(); builder.append::("A", 1).unwrap(); builder.append_null::("B").unwrap(); let expected_array = builder.build().unwrap(); diff --git a/arrow/src/ipc/reader.rs b/arrow/src/ipc/reader.rs index 8335fa4bc9e..086a5b15957 100644 --- a/arrow/src/ipc/reader.rs +++ b/arrow/src/ipc/reader.rs @@ -1627,7 +1627,7 @@ mod tests { let array1 = StringArray::from(vec!["foo", "bar", "baz"]); let array2 = BooleanArray::from(vec![true, false, true]); - let mut union_builder = UnionBuilder::new_dense(3); + let mut union_builder = UnionBuilder::new_dense(); union_builder.append::("a", 1).unwrap(); union_builder.append::("b", 10.1).unwrap(); union_builder.append_null::("b").unwrap(); @@ -1876,12 +1876,12 @@ mod tests { #[test] fn test_roundtrip_dense_union() { - check_union_with_builder(UnionBuilder::new_dense(6)); + check_union_with_builder(UnionBuilder::new_dense()); } #[test] fn test_roundtrip_sparse_union() { - check_union_with_builder(UnionBuilder::new_sparse(6)); + check_union_with_builder(UnionBuilder::new_sparse()); } /// Read gzipped JSON file diff --git a/arrow/src/ipc/writer.rs b/arrow/src/ipc/writer.rs index c817cb77c13..826e49ca87d 100644 --- a/arrow/src/ipc/writer.rs +++ b/arrow/src/ipc/writer.rs @@ -2040,7 +2040,7 @@ mod tests { ), true, )]); - let mut builder = UnionBuilder::new_sparse(5); + let mut builder = UnionBuilder::with_capacity_sparse(5); builder.append::("a", 1).unwrap(); builder.append_null::("a").unwrap(); builder.append::("c", 3.0).unwrap(); diff --git a/arrow/src/util/pretty.rs b/arrow/src/util/pretty.rs index 6f4d9e34a99..b92e658f1f2 100644 --- a/arrow/src/util/pretty.rs +++ b/arrow/src/util/pretty.rs @@ -649,7 +649,7 @@ mod tests { #[test] fn test_pretty_format_dense_union() -> Result<()> { - let mut builder = UnionBuilder::new_dense(4); + let mut builder = UnionBuilder::new_dense(); builder.append::("a", 1).unwrap(); builder.append::("b", 3.2234).unwrap(); builder.append_null::("b").unwrap(); @@ -690,7 +690,7 @@ mod tests { #[test] fn test_pretty_format_sparse_union() -> Result<()> { - let mut builder = UnionBuilder::new_sparse(4); + let mut builder = UnionBuilder::new_sparse(); builder.append::("a", 1).unwrap(); builder.append::("b", 3.2234).unwrap(); builder.append_null::("b").unwrap(); @@ -732,7 +732,7 @@ mod tests { #[test] fn test_pretty_format_nested_union() -> Result<()> { //Inner UnionArray - let mut builder = UnionBuilder::new_dense(5); + let mut builder = UnionBuilder::new_dense(); builder.append::("b", 1).unwrap(); builder.append::("c", 3.2234).unwrap(); builder.append_null::("c").unwrap();