1154 lines
34 KiB
Rust
1154 lines
34 KiB
Rust
// Licensed to the Apache Software Foundation (ASF) under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you under the Apache License, Version 2.0 (the
|
|
// "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing,
|
|
// software distributed under the License is distributed on an
|
|
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
// KIND, either express or implied. See the License for the
|
|
// specific language governing permissions and limitations
|
|
// under the License.
|
|
|
|
use arrow::array::{
|
|
Array, ArrayRef, BooleanArray, Decimal128Array, DictionaryArray, FixedSizeBinaryArray,
|
|
FixedSizeListBuilder, Int16Array, Int32Array, Int64Array, Int64Builder, ListArray, ListBuilder,
|
|
MapBuilder, NullArray, StringArray, StringBuilder, StringDictionaryBuilder, StructArray,
|
|
UInt8Array, UInt16Array, UInt16Builder, UnionArray,
|
|
};
|
|
use arrow::datatypes::Int16Type;
|
|
use arrow_array::StringViewArray;
|
|
use arrow_buffer::{Buffer, ScalarBuffer};
|
|
use arrow_data::ArrayData;
|
|
use arrow_data::transform::MutableArrayData;
|
|
use arrow_schema::{DataType, Field, Fields, UnionFields};
|
|
use std::sync::Arc;
|
|
|
|
#[allow(unused)]
|
|
fn create_decimal_array(array: Vec<Option<i128>>, precision: u8, scale: i8) -> Decimal128Array {
|
|
array
|
|
.into_iter()
|
|
.collect::<Decimal128Array>()
|
|
.with_precision_and_scale(precision, scale)
|
|
.unwrap()
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(not(feature = "force_validate"))]
|
|
fn test_decimal() {
|
|
let decimal_array =
|
|
create_decimal_array(vec![Some(1), Some(2), None, Some(3)], 10, 3).into_data();
|
|
let arrays = vec![&decimal_array];
|
|
let mut a = MutableArrayData::new(arrays, true, 3);
|
|
a.extend(0, 0, 3);
|
|
a.extend(0, 2, 3);
|
|
let result = a.freeze();
|
|
let array = Decimal128Array::from(result);
|
|
let expected = create_decimal_array(vec![Some(1), Some(2), None, None], 10, 3);
|
|
assert_eq!(array, expected);
|
|
}
|
|
#[test]
|
|
#[cfg(not(feature = "force_validate"))]
|
|
fn test_decimal_offset() {
|
|
let decimal_array = create_decimal_array(vec![Some(1), Some(2), None, Some(3)], 10, 3);
|
|
let decimal_array = decimal_array.slice(1, 3).into_data(); // 2, null, 3
|
|
let arrays = vec![&decimal_array];
|
|
let mut a = MutableArrayData::new(arrays, true, 2);
|
|
a.extend(0, 0, 2); // 2, null
|
|
let result = a.freeze();
|
|
let array = Decimal128Array::from(result);
|
|
let expected = create_decimal_array(vec![Some(2), None], 10, 3);
|
|
assert_eq!(array, expected);
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(not(feature = "force_validate"))]
|
|
fn test_decimal_null_offset_nulls() {
|
|
let decimal_array = create_decimal_array(vec![Some(1), Some(2), None, Some(3)], 10, 3);
|
|
let decimal_array = decimal_array.slice(1, 3).into_data(); // 2, null, 3
|
|
let arrays = vec![&decimal_array];
|
|
let mut a = MutableArrayData::new(arrays, true, 2);
|
|
a.extend(0, 0, 2); // 2, null
|
|
a.extend_nulls(3); // 2, null, null, null, null
|
|
a.extend(0, 1, 3); //2, null, null, null, null, null, 3
|
|
let result = a.freeze();
|
|
let array = Decimal128Array::from(result);
|
|
let expected =
|
|
create_decimal_array(vec![Some(2), None, None, None, None, None, Some(3)], 10, 3);
|
|
assert_eq!(array, expected);
|
|
}
|
|
|
|
/// tests extending from a primitive array w/ offset nor nulls
|
|
#[test]
|
|
fn test_primitive() {
|
|
let b = UInt8Array::from(vec![Some(1), Some(2), Some(3)]).into_data();
|
|
let arrays = vec![&b];
|
|
let mut a = MutableArrayData::new(arrays, false, 3);
|
|
a.extend(0, 0, 2);
|
|
let result = a.freeze();
|
|
let array = UInt8Array::from(result);
|
|
let expected = UInt8Array::from(vec![Some(1), Some(2)]);
|
|
assert_eq!(array, expected);
|
|
}
|
|
|
|
/// tests extending from a primitive array with offset w/ nulls
|
|
#[test]
|
|
fn test_primitive_offset() {
|
|
let b = UInt8Array::from(vec![Some(1), Some(2), Some(3)]).into_data();
|
|
let b = b.slice(1, 2);
|
|
let arrays = vec![&b];
|
|
let mut a = MutableArrayData::new(arrays, false, 2);
|
|
a.extend(0, 0, 2);
|
|
let result = a.freeze();
|
|
let array = UInt8Array::from(result);
|
|
let expected = UInt8Array::from(vec![Some(2), Some(3)]);
|
|
assert_eq!(array, expected);
|
|
}
|
|
|
|
/// tests extending from a primitive array with offset and nulls
|
|
#[test]
|
|
fn test_primitive_null_offset() {
|
|
let b = UInt8Array::from(vec![Some(1), None, Some(3)]);
|
|
let b = b.slice(1, 2).into_data();
|
|
let arrays = vec![&b];
|
|
let mut a = MutableArrayData::new(arrays, false, 2);
|
|
a.extend(0, 0, 2);
|
|
let result = a.freeze();
|
|
let array = UInt8Array::from(result);
|
|
let expected = UInt8Array::from(vec![None, Some(3)]);
|
|
assert_eq!(array, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_primitive_null_offset_nulls() {
|
|
let b = UInt8Array::from(vec![Some(1), Some(2), Some(3)]).into_data();
|
|
let b = b.slice(1, 2);
|
|
let arrays = vec![&b];
|
|
let mut a = MutableArrayData::new(arrays, true, 2);
|
|
a.extend(0, 0, 2);
|
|
a.extend_nulls(3);
|
|
a.extend(0, 1, 2);
|
|
let result = a.freeze();
|
|
let array = UInt8Array::from(result);
|
|
let expected = UInt8Array::from(vec![Some(2), Some(3), None, None, None, Some(3)]);
|
|
assert_eq!(array, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_list_null_offset() {
|
|
let int_builder = Int64Builder::with_capacity(24);
|
|
let mut builder = ListBuilder::<Int64Builder>::new(int_builder);
|
|
builder.values().append_slice(&[1, 2, 3]);
|
|
builder.append(true);
|
|
builder.values().append_slice(&[4, 5]);
|
|
builder.append(true);
|
|
builder.values().append_slice(&[6, 7, 8]);
|
|
builder.append(true);
|
|
let array = builder.finish().into_data();
|
|
let arrays = vec![&array];
|
|
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
mutable.extend(0, 0, 1);
|
|
|
|
let result = mutable.freeze();
|
|
let array = ListArray::from(result);
|
|
|
|
let int_builder = Int64Builder::with_capacity(24);
|
|
let mut builder = ListBuilder::<Int64Builder>::new(int_builder);
|
|
builder.values().append_slice(&[1, 2, 3]);
|
|
builder.append(true);
|
|
let expected = builder.finish();
|
|
|
|
assert_eq!(array, expected);
|
|
}
|
|
|
|
/// tests extending from a variable-sized (strings and binary) array w/ offset with nulls
|
|
#[test]
|
|
fn test_variable_sized_nulls() {
|
|
let array = StringArray::from(vec![Some("a"), Some("bc"), None, Some("defh")]).into_data();
|
|
let arrays = vec![&array];
|
|
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
mutable.extend(0, 1, 3);
|
|
|
|
let result = mutable.freeze();
|
|
let result = StringArray::from(result);
|
|
|
|
let expected = StringArray::from(vec![Some("bc"), None]);
|
|
assert_eq!(result, expected);
|
|
}
|
|
|
|
/// tests extending from a variable-sized (strings and binary) array
|
|
/// with an offset and nulls
|
|
#[test]
|
|
fn test_variable_sized_offsets() {
|
|
let array = StringArray::from(vec![Some("a"), Some("bc"), None, Some("defh")]);
|
|
let array = array.into_data().slice(1, 3);
|
|
|
|
let arrays = vec![&array];
|
|
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
mutable.extend(0, 0, 3);
|
|
|
|
let result = mutable.freeze();
|
|
let result = StringArray::from(result);
|
|
|
|
let expected = StringArray::from(vec![Some("bc"), None, Some("defh")]);
|
|
assert_eq!(result, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_string_offsets() {
|
|
let array = StringArray::from(vec![Some("a"), Some("bc"), None, Some("defh")]);
|
|
let array = array.into_data().slice(1, 3);
|
|
|
|
let arrays = vec![&array];
|
|
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
mutable.extend(0, 0, 3);
|
|
|
|
let result = mutable.freeze();
|
|
let result = StringArray::from(result);
|
|
|
|
let expected = StringArray::from(vec![Some("bc"), None, Some("defh")]);
|
|
assert_eq!(result, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_multiple_with_nulls() {
|
|
let array1 = StringArray::from(vec!["hello", "world"]).into_data();
|
|
let array2 = StringArray::from(vec![Some("1"), None]).into_data();
|
|
|
|
let arrays = vec![&array1, &array2];
|
|
|
|
let mut mutable = MutableArrayData::new(arrays, false, 5);
|
|
|
|
mutable.extend(0, 0, 2);
|
|
mutable.extend(1, 0, 2);
|
|
|
|
let result = mutable.freeze();
|
|
let result = StringArray::from(result);
|
|
|
|
let expected = StringArray::from(vec![Some("hello"), Some("world"), Some("1"), None]);
|
|
assert_eq!(result, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_string_null_offset_nulls() {
|
|
let array = StringArray::from(vec![Some("a"), Some("bc"), None, Some("defh")]);
|
|
let array = array.into_data().slice(1, 3);
|
|
|
|
let arrays = vec![&array];
|
|
|
|
let mut mutable = MutableArrayData::new(arrays, true, 0);
|
|
|
|
mutable.extend(0, 1, 3);
|
|
mutable.extend_nulls(1);
|
|
|
|
let result = mutable.freeze();
|
|
let result = StringArray::from(result);
|
|
|
|
let expected = StringArray::from(vec![None, Some("defh"), None]);
|
|
assert_eq!(result, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bool() {
|
|
let array = BooleanArray::from(vec![Some(false), Some(true), None, Some(false)]).into_data();
|
|
let arrays = vec![&array];
|
|
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
mutable.extend(0, 1, 3);
|
|
|
|
let result = mutable.freeze();
|
|
let result = BooleanArray::from(result);
|
|
|
|
let expected = BooleanArray::from(vec![Some(true), None]);
|
|
assert_eq!(result, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_null() {
|
|
let array1 = NullArray::new(10).into_data();
|
|
let array2 = NullArray::new(5).into_data();
|
|
let arrays = vec![&array1, &array2];
|
|
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
mutable.extend(0, 1, 3);
|
|
mutable.extend(1, 0, 1);
|
|
|
|
let result = mutable.freeze();
|
|
let result = NullArray::from(result);
|
|
|
|
let expected = NullArray::new(3);
|
|
assert_eq!(result, expected);
|
|
}
|
|
|
|
fn create_dictionary_array(values: &[&str], keys: &[Option<&str>]) -> ArrayData {
|
|
let values = StringArray::from(values.to_vec());
|
|
let mut builder =
|
|
StringDictionaryBuilder::<Int16Type>::new_with_dictionary(keys.len(), &values).unwrap();
|
|
for key in keys {
|
|
if let Some(v) = key {
|
|
builder.append(v).unwrap();
|
|
} else {
|
|
builder.append_null()
|
|
}
|
|
}
|
|
builder.finish().into_data()
|
|
}
|
|
|
|
#[test]
|
|
fn test_dictionary() {
|
|
// (a, b, c), (0, 1, 0, 2) => (a, b, a, c)
|
|
let array = create_dictionary_array(&["a", "b", "c"], &[Some("a"), Some("b"), None, Some("c")]);
|
|
let arrays = vec![&array];
|
|
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
mutable.extend(0, 1, 3);
|
|
|
|
let result = mutable.freeze();
|
|
let result = DictionaryArray::from(result);
|
|
|
|
let expected = Int16Array::from(vec![Some(1), None]);
|
|
assert_eq!(result.keys(), &expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_struct() {
|
|
let strings: ArrayRef = Arc::new(StringArray::from(vec![
|
|
Some("joe"),
|
|
None,
|
|
None,
|
|
Some("mark"),
|
|
Some("doe"),
|
|
]));
|
|
let ints: ArrayRef = Arc::new(Int32Array::from(vec![
|
|
Some(1),
|
|
Some(2),
|
|
Some(3),
|
|
Some(4),
|
|
Some(5),
|
|
]));
|
|
|
|
let array = StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
|
|
.unwrap()
|
|
.into_data();
|
|
let arrays = vec![&array];
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
mutable.extend(0, 1, 3);
|
|
let data = mutable.freeze();
|
|
let array = StructArray::from(data);
|
|
|
|
let expected =
|
|
StructArray::try_from(vec![("f1", strings.slice(1, 2)), ("f2", ints.slice(1, 2))]).unwrap();
|
|
assert_eq!(array, expected)
|
|
}
|
|
|
|
#[test]
|
|
fn test_struct_offset() {
|
|
let strings: ArrayRef = Arc::new(StringArray::from(vec![
|
|
Some("joe"),
|
|
None,
|
|
None,
|
|
Some("mark"),
|
|
Some("doe"),
|
|
]));
|
|
let ints: ArrayRef = Arc::new(Int32Array::from(vec![
|
|
Some(1),
|
|
Some(2),
|
|
Some(3),
|
|
Some(4),
|
|
Some(5),
|
|
]));
|
|
|
|
let array = StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
|
|
.unwrap()
|
|
.into_data()
|
|
.slice(1, 3);
|
|
let arrays = vec![&array];
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
mutable.extend(0, 1, 3);
|
|
let data = mutable.freeze();
|
|
let array = StructArray::from(data);
|
|
|
|
let expected_strings: ArrayRef = Arc::new(StringArray::from(vec![None, Some("mark")]));
|
|
let expected =
|
|
StructArray::try_from(vec![("f1", expected_strings), ("f2", ints.slice(2, 2))]).unwrap();
|
|
|
|
assert_eq!(array, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_struct_nulls() {
|
|
let strings: ArrayRef = Arc::new(StringArray::from(vec![
|
|
Some("joe"),
|
|
None,
|
|
None,
|
|
Some("mark"),
|
|
Some("doe"),
|
|
]));
|
|
let ints: ArrayRef = Arc::new(Int32Array::from(vec![
|
|
Some(1),
|
|
Some(2),
|
|
None,
|
|
Some(4),
|
|
Some(5),
|
|
]));
|
|
|
|
let array = StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
|
|
.unwrap()
|
|
.into_data();
|
|
let arrays = vec![&array];
|
|
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
mutable.extend(0, 1, 3);
|
|
let data = mutable.freeze();
|
|
let array = StructArray::from(data);
|
|
|
|
let v: Vec<Option<&str>> = vec![None, None];
|
|
let expected_string = Arc::new(StringArray::from(v)) as ArrayRef;
|
|
let expected_int = Arc::new(Int32Array::from(vec![Some(2), None])) as ArrayRef;
|
|
|
|
let expected =
|
|
StructArray::try_from(vec![("f1", expected_string), ("f2", expected_int)]).unwrap();
|
|
assert_eq!(array, expected)
|
|
}
|
|
|
|
#[test]
|
|
fn test_struct_many() {
|
|
let strings: ArrayRef = Arc::new(StringArray::from(vec![
|
|
Some("joe"),
|
|
None,
|
|
None,
|
|
Some("mark"),
|
|
Some("doe"),
|
|
]));
|
|
let ints: ArrayRef = Arc::new(Int32Array::from(vec![
|
|
Some(1),
|
|
Some(2),
|
|
None,
|
|
Some(4),
|
|
Some(5),
|
|
]));
|
|
|
|
let array = StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())])
|
|
.unwrap()
|
|
.into_data();
|
|
let arrays = vec![&array, &array];
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
mutable.extend(0, 1, 3);
|
|
mutable.extend(1, 0, 2);
|
|
let data = mutable.freeze();
|
|
let array = StructArray::from(data);
|
|
|
|
let expected_string =
|
|
Arc::new(StringArray::from(vec![None, None, Some("joe"), None])) as ArrayRef;
|
|
let expected_int =
|
|
Arc::new(Int32Array::from(vec![Some(2), None, Some(1), Some(2)])) as ArrayRef;
|
|
|
|
let expected =
|
|
StructArray::try_from(vec![("f1", expected_string), ("f2", expected_int)]).unwrap();
|
|
assert_eq!(array, expected)
|
|
}
|
|
|
|
#[test]
|
|
fn test_union_dense() {
|
|
// Input data
|
|
let strings: ArrayRef = Arc::new(StringArray::from(vec![
|
|
Some("joe"),
|
|
Some("mark"),
|
|
Some("doe"),
|
|
]));
|
|
let ints: ArrayRef = Arc::new(Int32Array::from(vec![
|
|
Some(1),
|
|
Some(2),
|
|
Some(3),
|
|
Some(4),
|
|
Some(5),
|
|
]));
|
|
let offsets = [0, 0, 1, 1, 2, 2, 3, 4]
|
|
.into_iter()
|
|
.collect::<ScalarBuffer<i32>>();
|
|
let type_ids = [42, 84, 42, 84, 84, 42, 84, 84]
|
|
.into_iter()
|
|
.collect::<ScalarBuffer<i8>>();
|
|
|
|
let union_fields = [
|
|
(84, Arc::new(Field::new("int", DataType::Int32, false))),
|
|
(42, Arc::new(Field::new("string", DataType::Utf8, false))),
|
|
]
|
|
.into_iter()
|
|
.collect::<UnionFields>();
|
|
|
|
let array = UnionArray::try_new(
|
|
union_fields.clone(),
|
|
type_ids,
|
|
Some(offsets),
|
|
vec![ints, strings],
|
|
)
|
|
.unwrap()
|
|
.into_data();
|
|
let arrays = vec![&array];
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
// Slice it by `MutableArrayData`
|
|
mutable.extend(0, 4, 7);
|
|
let data = mutable.freeze();
|
|
let array = UnionArray::from(data);
|
|
|
|
// Expected data
|
|
let strings: ArrayRef = Arc::new(StringArray::from(vec![Some("doe")]));
|
|
let ints: ArrayRef = Arc::new(Int32Array::from(vec![Some(3), Some(4)]));
|
|
let offsets = [0, 0, 1].into_iter().collect::<ScalarBuffer<i32>>();
|
|
let type_ids = [84, 42, 84].into_iter().collect::<ScalarBuffer<i8>>();
|
|
|
|
let expected =
|
|
UnionArray::try_new(union_fields, type_ids, Some(offsets), vec![ints, strings]).unwrap();
|
|
|
|
assert_eq!(array.to_data(), expected.to_data());
|
|
}
|
|
|
|
#[test]
|
|
fn test_binary_fixed_sized_offsets() {
|
|
let array =
|
|
FixedSizeBinaryArray::try_from_iter(vec![vec![0, 0], vec![0, 1], vec![0, 2]].into_iter())
|
|
.expect("Failed to create FixedSizeBinaryArray from iterable");
|
|
let array = array.slice(1, 2).into_data();
|
|
// = [[0, 1], [0, 2]] due to the offset = 1
|
|
|
|
let arrays = vec![&array];
|
|
|
|
let mut mutable = MutableArrayData::new(arrays, false, 0);
|
|
|
|
mutable.extend(0, 1, 2);
|
|
mutable.extend(0, 0, 1);
|
|
|
|
let result = mutable.freeze();
|
|
let result = FixedSizeBinaryArray::from(result);
|
|
|
|
let expected = FixedSizeBinaryArray::try_from_iter(vec![vec![0, 2], vec![0, 1]].into_iter())
|
|
.expect("Failed to create FixedSizeBinaryArray from iterable");
|
|
assert_eq!(result, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_list_append() {
|
|
let mut builder = ListBuilder::<Int64Builder>::new(Int64Builder::with_capacity(24));
|
|
builder.values().append_slice(&[1, 2, 3]);
|
|
builder.append(true);
|
|
builder.values().append_slice(&[4, 5]);
|
|
builder.append(true);
|
|
builder.values().append_slice(&[6, 7, 8]);
|
|
builder.values().append_slice(&[9, 10, 11]);
|
|
builder.append(true);
|
|
let a = builder.finish().into_data();
|
|
|
|
let a_builder = Int64Builder::with_capacity(24);
|
|
let mut a_builder = ListBuilder::<Int64Builder>::new(a_builder);
|
|
a_builder.values().append_slice(&[12, 13]);
|
|
a_builder.append(true);
|
|
a_builder.append(true);
|
|
a_builder.values().append_slice(&[14, 15]);
|
|
a_builder.append(true);
|
|
let b = a_builder.finish().into_data();
|
|
|
|
let c = b.slice(1, 2);
|
|
|
|
let mut mutable = MutableArrayData::new(vec![&a, &b, &c], false, 1);
|
|
mutable.extend(0, 0, a.len());
|
|
mutable.extend(1, 0, b.len());
|
|
mutable.extend(2, 0, c.len());
|
|
|
|
let finished = mutable.freeze();
|
|
|
|
let expected_int_array = Int64Array::from(vec![
|
|
Some(1),
|
|
Some(2),
|
|
Some(3),
|
|
Some(4),
|
|
Some(5),
|
|
Some(6),
|
|
Some(7),
|
|
Some(8),
|
|
Some(9),
|
|
Some(10),
|
|
Some(11),
|
|
// append first array
|
|
Some(12),
|
|
Some(13),
|
|
Some(14),
|
|
Some(15),
|
|
// append second array
|
|
Some(14),
|
|
Some(15),
|
|
]);
|
|
let list_value_offsets = Buffer::from_slice_ref([0i32, 3, 5, 11, 13, 13, 15, 15, 17]);
|
|
let expected_list_data = ArrayData::try_new(
|
|
DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))),
|
|
8,
|
|
None,
|
|
0,
|
|
vec![list_value_offsets],
|
|
vec![expected_int_array.into_data()],
|
|
)
|
|
.unwrap();
|
|
assert_eq!(finished, expected_list_data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_list_nulls_append() {
|
|
let mut builder = ListBuilder::<Int64Builder>::new(Int64Builder::with_capacity(32));
|
|
builder.values().append_slice(&[1, 2, 3]);
|
|
builder.append(true);
|
|
builder.values().append_slice(&[4, 5]);
|
|
builder.append(true);
|
|
builder.append(false);
|
|
builder.values().append_slice(&[6, 7, 8]);
|
|
builder.values().append_null();
|
|
builder.values().append_null();
|
|
builder.values().append_slice(&[9, 10, 11]);
|
|
builder.append(true);
|
|
let a = builder.finish().into_data();
|
|
|
|
let mut builder = ListBuilder::<Int64Builder>::new(Int64Builder::with_capacity(32));
|
|
builder.values().append_slice(&[12, 13]);
|
|
builder.append(true);
|
|
builder.append(false);
|
|
builder.append(true);
|
|
builder.values().append_null();
|
|
builder.values().append_null();
|
|
builder.values().append_slice(&[14, 15]);
|
|
builder.append(true);
|
|
let b = builder.finish().into_data();
|
|
let c = b.slice(1, 2);
|
|
let d = b.slice(2, 2);
|
|
|
|
let mut mutable = MutableArrayData::new(vec![&a, &b, &c, &d], false, 10);
|
|
|
|
mutable.extend(0, 0, a.len());
|
|
mutable.extend(1, 0, b.len());
|
|
mutable.extend(2, 0, c.len());
|
|
mutable.extend(3, 0, d.len());
|
|
let result = mutable.freeze();
|
|
|
|
let expected_int_array = Int64Array::from(vec![
|
|
Some(1),
|
|
Some(2),
|
|
Some(3),
|
|
Some(4),
|
|
Some(5),
|
|
Some(6),
|
|
Some(7),
|
|
Some(8),
|
|
None,
|
|
None,
|
|
Some(9),
|
|
Some(10),
|
|
Some(11),
|
|
// second array
|
|
Some(12),
|
|
Some(13),
|
|
None,
|
|
None,
|
|
Some(14),
|
|
Some(15),
|
|
// slice(1, 2) results in no values added
|
|
None,
|
|
None,
|
|
Some(14),
|
|
Some(15),
|
|
]);
|
|
let list_value_offsets =
|
|
Buffer::from_slice_ref([0, 3, 5, 5, 13, 15, 15, 15, 19, 19, 19, 19, 23]);
|
|
let expected_list_data = ArrayData::try_new(
|
|
DataType::List(Arc::new(Field::new_list_field(DataType::Int64, true))),
|
|
12,
|
|
Some(Buffer::from(&[0b11011011, 0b1110])),
|
|
0,
|
|
vec![list_value_offsets],
|
|
vec![expected_int_array.into_data()],
|
|
)
|
|
.unwrap();
|
|
assert_eq!(result, expected_list_data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_map_nulls_append() {
|
|
let mut builder = MapBuilder::<Int64Builder, Int64Builder>::new(
|
|
None,
|
|
Int64Builder::with_capacity(32),
|
|
Int64Builder::with_capacity(32),
|
|
);
|
|
builder.keys().append_slice(&[1, 2, 3]);
|
|
builder.values().append_slice(&[1, 2, 3]);
|
|
builder.append(true).unwrap();
|
|
builder.keys().append_slice(&[4, 5]);
|
|
builder.values().append_slice(&[4, 5]);
|
|
builder.append(true).unwrap();
|
|
builder.append(false).unwrap();
|
|
builder.keys().append_slice(&[6, 7, 8, 100, 101, 9, 10, 11]);
|
|
builder.values().append_slice(&[6, 7, 8]);
|
|
builder.values().append_null();
|
|
builder.values().append_null();
|
|
builder.values().append_slice(&[9, 10, 11]);
|
|
builder.append(true).unwrap();
|
|
|
|
let a = builder.finish().into_data();
|
|
|
|
let mut builder = MapBuilder::<Int64Builder, Int64Builder>::new(
|
|
None,
|
|
Int64Builder::with_capacity(32),
|
|
Int64Builder::with_capacity(32),
|
|
);
|
|
|
|
builder.keys().append_slice(&[12, 13]);
|
|
builder.values().append_slice(&[12, 13]);
|
|
builder.append(true).unwrap();
|
|
builder.append(false).unwrap();
|
|
builder.append(true).unwrap();
|
|
builder.keys().append_slice(&[100, 101, 14, 15]);
|
|
builder.values().append_null();
|
|
builder.values().append_null();
|
|
builder.values().append_slice(&[14, 15]);
|
|
builder.append(true).unwrap();
|
|
|
|
let b = builder.finish().into_data();
|
|
let c = b.slice(1, 2);
|
|
let d = b.slice(2, 2);
|
|
|
|
let mut mutable = MutableArrayData::new(vec![&a, &b, &c, &d], false, 10);
|
|
|
|
mutable.extend(0, 0, a.len());
|
|
mutable.extend(1, 0, b.len());
|
|
mutable.extend(2, 0, c.len());
|
|
mutable.extend(3, 0, d.len());
|
|
let result = mutable.freeze();
|
|
|
|
let expected_key_array = Int64Array::from(vec![
|
|
Some(1),
|
|
Some(2),
|
|
Some(3),
|
|
Some(4),
|
|
Some(5),
|
|
Some(6),
|
|
Some(7),
|
|
Some(8),
|
|
Some(100),
|
|
Some(101),
|
|
Some(9),
|
|
Some(10),
|
|
Some(11),
|
|
// second array
|
|
Some(12),
|
|
Some(13),
|
|
Some(100),
|
|
Some(101),
|
|
Some(14),
|
|
Some(15),
|
|
// slice(1, 2) results in no values added
|
|
Some(100),
|
|
Some(101),
|
|
Some(14),
|
|
Some(15),
|
|
]);
|
|
|
|
let expected_value_array = Int64Array::from(vec![
|
|
Some(1),
|
|
Some(2),
|
|
Some(3),
|
|
Some(4),
|
|
Some(5),
|
|
Some(6),
|
|
Some(7),
|
|
Some(8),
|
|
None,
|
|
None,
|
|
Some(9),
|
|
Some(10),
|
|
Some(11),
|
|
// second array
|
|
Some(12),
|
|
Some(13),
|
|
None,
|
|
None,
|
|
Some(14),
|
|
Some(15),
|
|
// slice(1, 2) results in no values added
|
|
None,
|
|
None,
|
|
Some(14),
|
|
Some(15),
|
|
]);
|
|
|
|
let expected_entry_array = StructArray::from(vec![
|
|
(
|
|
Arc::new(Field::new("keys", DataType::Int64, false)),
|
|
Arc::new(expected_key_array) as ArrayRef,
|
|
),
|
|
(
|
|
Arc::new(Field::new("values", DataType::Int64, true)),
|
|
Arc::new(expected_value_array) as ArrayRef,
|
|
),
|
|
]);
|
|
|
|
let map_offsets = Buffer::from_slice_ref([0, 3, 5, 5, 13, 15, 15, 15, 19, 19, 19, 19, 23]);
|
|
|
|
let expected_list_data = ArrayData::try_new(
|
|
DataType::Map(
|
|
Arc::new(Field::new(
|
|
"entries",
|
|
DataType::Struct(Fields::from(vec![
|
|
Field::new("keys", DataType::Int64, false),
|
|
Field::new("values", DataType::Int64, true),
|
|
])),
|
|
false,
|
|
)),
|
|
false,
|
|
),
|
|
12,
|
|
Some(Buffer::from(&[0b11011011, 0b1110])),
|
|
0,
|
|
vec![map_offsets],
|
|
vec![expected_entry_array.into_data()],
|
|
)
|
|
.unwrap();
|
|
assert_eq!(result, expected_list_data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_map_keys_values_append() {
|
|
let mut builder = MapBuilder::<Int64Builder, Int64Builder>::new(
|
|
None,
|
|
Int64Builder::with_capacity(32),
|
|
Int64Builder::with_capacity(32),
|
|
);
|
|
let (keys, values) = builder.entries();
|
|
keys.append_slice(&[1, 2, 3]);
|
|
values.append_slice(&[1, 3, 4]);
|
|
builder.append(true).unwrap();
|
|
|
|
let (keys, values) = builder.entries();
|
|
keys.append_slice(&[4, 5]);
|
|
values.append_slice(&[4, 6]);
|
|
builder.append(true).unwrap();
|
|
|
|
builder.append(false).unwrap();
|
|
|
|
let map = builder.finish();
|
|
assert!(map.is_null(2));
|
|
|
|
let first = map.value(0);
|
|
let keys = first
|
|
.column(0)
|
|
.as_any()
|
|
.downcast_ref::<Int64Array>()
|
|
.unwrap();
|
|
let values = first
|
|
.column(1)
|
|
.as_any()
|
|
.downcast_ref::<Int64Array>()
|
|
.unwrap();
|
|
assert_eq!(keys, &Int64Array::from(vec![Some(1), Some(2), Some(3)]));
|
|
assert_eq!(values, &Int64Array::from(vec![Some(1), Some(3), Some(4)]));
|
|
|
|
let second = map.value(1);
|
|
let keys = second
|
|
.column(0)
|
|
.as_any()
|
|
.downcast_ref::<Int64Array>()
|
|
.unwrap();
|
|
let values = second
|
|
.column(1)
|
|
.as_any()
|
|
.downcast_ref::<Int64Array>()
|
|
.unwrap();
|
|
assert_eq!(keys, &Int64Array::from(vec![Some(4), Some(5)]));
|
|
assert_eq!(values, &Int64Array::from(vec![Some(4), Some(6)]));
|
|
}
|
|
|
|
#[test]
|
|
fn test_list_of_strings_append() {
|
|
// [["alpha", "beta", None]]
|
|
let mut builder = ListBuilder::new(StringBuilder::new());
|
|
builder.values().append_value("Hello");
|
|
builder.values().append_value("Arrow");
|
|
builder.values().append_null();
|
|
builder.append(true);
|
|
let a = builder.finish().into_data();
|
|
|
|
// [["alpha", "beta"], [None], ["gamma", "delta", None]]
|
|
let mut builder = ListBuilder::new(StringBuilder::new());
|
|
builder.values().append_value("alpha");
|
|
builder.values().append_value("beta");
|
|
builder.append(true);
|
|
builder.values().append_null();
|
|
builder.append(true);
|
|
builder.values().append_value("gamma");
|
|
builder.values().append_value("delta");
|
|
builder.values().append_null();
|
|
builder.append(true);
|
|
let b = builder.finish().into_data();
|
|
|
|
let mut mutable = MutableArrayData::new(vec![&a, &b], false, 10);
|
|
|
|
mutable.extend(0, 0, a.len());
|
|
mutable.extend(1, 0, b.len());
|
|
mutable.extend(1, 1, 3);
|
|
mutable.extend(1, 0, 0);
|
|
let result = mutable.freeze();
|
|
|
|
let expected_string_array = StringArray::from(vec![
|
|
// extend a[0..a.len()]
|
|
// a[0]
|
|
Some("Hello"),
|
|
Some("Arrow"),
|
|
None,
|
|
// extend b[0..b.len()]
|
|
// b[0]
|
|
Some("alpha"),
|
|
Some("beta"),
|
|
// b[1]
|
|
None,
|
|
// b[2]
|
|
Some("gamma"),
|
|
Some("delta"),
|
|
None,
|
|
// extend b[1..3]
|
|
// b[1]
|
|
None,
|
|
// b[2]
|
|
Some("gamma"),
|
|
Some("delta"),
|
|
None,
|
|
// extend b[0..0]
|
|
]);
|
|
let list_value_offsets = Buffer::from_slice_ref([0, 3, 5, 6, 9, 10, 13]);
|
|
let expected_list_data = ArrayData::try_new(
|
|
DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))),
|
|
6,
|
|
None,
|
|
0,
|
|
vec![list_value_offsets],
|
|
vec![expected_string_array.into_data()],
|
|
)
|
|
.unwrap();
|
|
assert_eq!(result, expected_list_data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_fixed_size_binary_append() {
|
|
let a = vec![Some(vec![1, 2]), Some(vec![3, 4]), Some(vec![5, 6])];
|
|
let a = FixedSizeBinaryArray::try_from_sparse_iter_with_size(a.into_iter(), 2)
|
|
.expect("Failed to create FixedSizeBinaryArray from iterable")
|
|
.into_data();
|
|
|
|
let b = vec![
|
|
None,
|
|
Some(vec![7, 8]),
|
|
Some(vec![9, 10]),
|
|
None,
|
|
Some(vec![13, 14]),
|
|
None,
|
|
];
|
|
let b = FixedSizeBinaryArray::try_from_sparse_iter_with_size(b.into_iter(), 2)
|
|
.expect("Failed to create FixedSizeBinaryArray from iterable")
|
|
.into_data();
|
|
|
|
let mut mutable = MutableArrayData::new(vec![&a, &b], false, 10);
|
|
|
|
mutable.extend(0, 0, a.len());
|
|
mutable.extend(1, 0, b.len());
|
|
mutable.extend(1, 1, 4);
|
|
mutable.extend(1, 2, 3);
|
|
mutable.extend(1, 5, 5);
|
|
let result = mutable.freeze();
|
|
|
|
let expected = vec![
|
|
// a
|
|
Some(vec![1, 2]),
|
|
Some(vec![3, 4]),
|
|
Some(vec![5, 6]),
|
|
// b
|
|
None,
|
|
Some(vec![7, 8]),
|
|
Some(vec![9, 10]),
|
|
None,
|
|
Some(vec![13, 14]),
|
|
None,
|
|
// b[1..4]
|
|
Some(vec![7, 8]),
|
|
Some(vec![9, 10]),
|
|
None,
|
|
// b[2..3]
|
|
Some(vec![9, 10]),
|
|
// b[4..4]
|
|
];
|
|
let expected = FixedSizeBinaryArray::try_from_sparse_iter_with_size(expected.into_iter(), 2)
|
|
.expect("Failed to create FixedSizeBinaryArray from iterable")
|
|
.into_data();
|
|
assert_eq!(result, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_extend_nulls() {
|
|
let int = Int32Array::from(vec![1, 2, 3, 4]).into_data();
|
|
let mut mutable = MutableArrayData::new(vec![&int], true, 4);
|
|
mutable.extend(0, 2, 3);
|
|
mutable.extend_nulls(2);
|
|
|
|
let data = mutable.freeze();
|
|
data.validate_full().unwrap();
|
|
let out = Int32Array::from(data);
|
|
|
|
assert_eq!(out.null_count(), 2);
|
|
assert_eq!(out.iter().collect::<Vec<_>>(), vec![Some(3), None, None]);
|
|
}
|
|
|
|
#[test]
|
|
#[should_panic(expected = "MutableArrayData not nullable")]
|
|
fn test_extend_nulls_panic() {
|
|
let int = Int32Array::from(vec![1, 2, 3, 4]).into_data();
|
|
let mut mutable = MutableArrayData::new(vec![&int], false, 4);
|
|
mutable.extend_nulls(2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_string_view() {
|
|
let a1 =
|
|
StringViewArray::from(vec!["foo", "very long string over 12 bytes", "bar"]).into_data();
|
|
let a2 = StringViewArray::from_iter(vec![
|
|
Some("bar"),
|
|
None,
|
|
Some("long string also over 12 bytes"),
|
|
])
|
|
.into_data();
|
|
|
|
a1.validate_full().unwrap();
|
|
a2.validate_full().unwrap();
|
|
|
|
let mut mutable = MutableArrayData::new(vec![&a1, &a2], false, 4);
|
|
mutable.extend(1, 0, 1);
|
|
mutable.extend(0, 1, 2);
|
|
mutable.extend(0, 0, 1);
|
|
mutable.extend(1, 2, 3);
|
|
|
|
let array = StringViewArray::from(mutable.freeze());
|
|
assert_eq!(array.data_buffers().len(), 2);
|
|
// Should have reused data buffers
|
|
assert_eq!(array.data_buffers()[0].as_ptr(), a1.buffers()[1].as_ptr());
|
|
assert_eq!(array.data_buffers()[1].as_ptr(), a2.buffers()[1].as_ptr());
|
|
|
|
let v = array.iter().collect::<Vec<_>>();
|
|
assert_eq!(
|
|
v,
|
|
vec![
|
|
Some("bar"),
|
|
Some("very long string over 12 bytes"),
|
|
Some("foo"),
|
|
Some("long string also over 12 bytes")
|
|
]
|
|
)
|
|
}
|
|
|
|
#[test]
|
|
#[should_panic(expected = "Arrays with inconsistent types passed to MutableArrayData")]
|
|
fn test_mixed_types() {
|
|
let a = StringArray::from(vec!["abc", "def"]).to_data();
|
|
let b = Int32Array::from(vec![1, 2, 3]).to_data();
|
|
MutableArrayData::new(vec![&a, &b], false, 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_fixed_size_list_append() {
|
|
let int_builder = UInt16Builder::with_capacity(64);
|
|
let mut builder = FixedSizeListBuilder::<UInt16Builder>::new(int_builder, 2);
|
|
builder.values().append_slice(&[1, 2]);
|
|
builder.append(true);
|
|
builder.values().append_slice(&[3, 4]);
|
|
builder.append(false);
|
|
builder.values().append_slice(&[5, 6]);
|
|
builder.append(true);
|
|
let a = builder.finish().into_data();
|
|
|
|
let a_builder = UInt16Builder::with_capacity(64);
|
|
let mut a_builder = FixedSizeListBuilder::<UInt16Builder>::new(a_builder, 2);
|
|
a_builder.values().append_slice(&[7, 8]);
|
|
a_builder.append(true);
|
|
a_builder.values().append_slice(&[9, 10]);
|
|
a_builder.append(true);
|
|
a_builder.values().append_slice(&[11, 12]);
|
|
a_builder.append(false);
|
|
a_builder.values().append_slice(&[13, 14]);
|
|
a_builder.append(true);
|
|
a_builder.values().append_null();
|
|
a_builder.values().append_null();
|
|
a_builder.append(true);
|
|
let b = a_builder.finish().into_data();
|
|
|
|
let mut mutable = MutableArrayData::new(vec![&a, &b], false, 10);
|
|
mutable.extend(0, 0, a.len());
|
|
mutable.extend(1, 0, b.len());
|
|
|
|
// append array
|
|
mutable.extend(1, 1, 4);
|
|
mutable.extend(1, 2, 3);
|
|
|
|
let finished = mutable.freeze();
|
|
|
|
let expected_int_array = UInt16Array::from(vec![
|
|
Some(1),
|
|
Some(2),
|
|
Some(3),
|
|
Some(4),
|
|
Some(5),
|
|
Some(6),
|
|
// append first array
|
|
Some(7),
|
|
Some(8),
|
|
Some(9),
|
|
Some(10),
|
|
Some(11),
|
|
Some(12),
|
|
Some(13),
|
|
Some(14),
|
|
None,
|
|
None,
|
|
// append slice(1, 3)
|
|
Some(9),
|
|
Some(10),
|
|
Some(11),
|
|
Some(12),
|
|
Some(13),
|
|
Some(14),
|
|
// append slice(2, 1)
|
|
Some(11),
|
|
Some(12),
|
|
]);
|
|
let expected_fixed_size_list_data = ArrayData::try_new(
|
|
DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::UInt16, true)), 2),
|
|
12,
|
|
Some(Buffer::from(&[0b11011101, 0b101])),
|
|
0,
|
|
vec![],
|
|
vec![expected_int_array.to_data()],
|
|
)
|
|
.unwrap();
|
|
assert_eq!(finished, expected_fixed_size_list_data);
|
|
}
|