You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2021/07/13 21:06:35 UTC
[GitHub] [arrow-rs] alamb commented on a change in pull request #491: Minimal MapArray support
alamb commented on a change in pull request #491:
URL: https://github.com/apache/arrow-rs/pull/491#discussion_r669078940
##########
File path: arrow/Cargo.toml
##########
@@ -58,13 +58,15 @@ multiversion = "0.6.1"
bitflags = "1.2.1"
[features]
-default = ["csv", "ipc"]
+default = ["csv", "ipc", "experimental-api"]
Review comment:
I recommend that `experimental-api` is not enabled by default as an additional safeguard against someone being surprised if a new arrow release breaks it
##########
File path: arrow/src/array/array_map.rs
##########
@@ -0,0 +1,380 @@
+// 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 std::any::Any;
+use std::fmt;
+use std::mem;
+
+use super::make_array;
+use super::{
+ array::print_long_array, raw_pointer::RawPtrBox, Array, ArrayData, ArrayRef,
+};
+use crate::datatypes::{ArrowNativeType, DataType};
+use crate::error::ArrowError;
+
+/// A nested array type where each record is a key-value map.
+/// Keys should always be non-null, but values can be null.
+///
+/// [MapArray] is physically a [ListArray] that has a [StructArray]
+/// with 2 child fields.
+pub struct MapArray {
+ data: ArrayData,
+ values: ArrayRef,
+ value_offsets: RawPtrBox<i32>,
+}
+
+impl MapArray {
+ #[cfg(feature = "experimental-api")]
+ /// Returns a reference to the keys of this map.
+ pub fn keys(&self) -> ArrayRef {
+ make_array(self.values.data().child_data()[0].clone())
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns a reference to the values of this map.
+ pub fn values(&self) -> ArrayRef {
+ make_array(self.values.data().child_data()[1].clone())
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns the data type of the map's keys.
+ pub fn key_type(&self) -> DataType {
+ self.values.data().child_data()[0].data_type().clone()
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns the data type of the map's values.
+ pub fn value_type(&self) -> DataType {
+ self.values.data().child_data()[1].data_type().clone()
+ }
+
+ /// Returns ith value of this map array.
+ /// # Safety
+ /// Caller must ensure that the index is within the array bounds
+ pub unsafe fn value_unchecked(&self, i: usize) -> ArrayRef {
+ let end = *self.value_offsets().get_unchecked(i + 1);
+ let start = *self.value_offsets().get_unchecked(i);
+ self.values
+ .slice(start.to_usize().unwrap(), (end - start).to_usize().unwrap())
+ }
+
+ /// Returns ith value of this map array.
+ pub fn value(&self, i: usize) -> ArrayRef {
+ let end = self.value_offsets()[i + 1] as usize;
+ let start = self.value_offsets()[i] as usize;
+ self.values.slice(start, end - start)
+ }
+
+ /// Returns the offset values in the offsets buffer
+ #[inline]
+ pub fn value_offsets(&self) -> &[i32] {
+ // Soundness
+ // pointer alignment & location is ensured by RawPtrBox
+ // buffer bounds/offset is ensured by the ArrayData instance.
+ unsafe {
+ std::slice::from_raw_parts(
+ self.value_offsets.as_ptr().add(self.data.offset()),
+ self.len() + 1,
+ )
+ }
+ }
+
+ /// Returns the length for value at index `i`.
+ #[inline]
+ pub fn value_length(&self, i: usize) -> i32 {
+ let offsets = self.value_offsets();
+ offsets[i + 1] - offsets[i]
+ }
+}
+
+impl From<ArrayData> for MapArray {
+ fn from(data: ArrayData) -> Self {
+ Self::try_new_from_array_data(data)
+ .expect("Expected infallable creation of MapArray from ArrayData failed")
+ }
+}
+
+impl MapArray {
+ fn try_new_from_array_data(data: ArrayData) -> Result<Self, ArrowError> {
+ if data.buffers().len() != 1 {
+ return Err(ArrowError::InvalidArgumentError(
+ format!("MapArray data should contain a single buffer only (value offsets), had {}",
+ data.len())));
+ }
+
+ if data.child_data().len() != 1 {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a single child array (values array), had {}",
+ data.child_data().len()
+ )));
+ }
+
+ let entries = data.child_data()[0].clone();
+
+ if let DataType::Struct(fields) = entries.data_type() {
+ if fields.len() != 2 {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a struct array with 2 fields, have {} fields",
+ fields.len()
+ )));
+ }
+ } else {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a struct array child, found {:?}",
+ entries.data_type()
+ )));
+ }
+
+ let values = make_array(entries);
+ let value_offsets = data.buffers()[0].as_ptr();
+
+ let value_offsets = unsafe { RawPtrBox::<i32>::new(value_offsets) };
+ unsafe {
+ if (*value_offsets.as_ptr().offset(0)) != 0 {
+ return Err(ArrowError::InvalidArgumentError(String::from(
+ "offsets do not start at zero",
+ )));
+ }
+ }
+ Ok(Self {
+ data,
+ values,
+ value_offsets,
+ })
+ }
+}
+
+impl Array for MapArray {
+ fn as_any(&self) -> &Any {
+ self
+ }
+
+ fn data(&self) -> &ArrayData {
+ &self.data
+ }
+
+ /// Returns the total number of bytes of memory occupied by the buffers owned by this [MapArray].
+ fn get_buffer_memory_size(&self) -> usize {
+ self.data.get_buffer_memory_size()
+ }
+
+ /// Returns the total number of bytes of memory occupied physically by this [MapArray].
+ fn get_array_memory_size(&self) -> usize {
+ self.data.get_array_memory_size() + mem::size_of_val(self)
+ }
+}
+
+impl fmt::Debug for MapArray {
+ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(f, "MapArray\n[\n")?;
+ print_long_array(self, f, |array, index, f| {
+ fmt::Debug::fmt(&array.value(index), f)
+ })?;
+ write!(f, "]")
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use crate::{
+ array::ArrayData,
+ array::{Int32Array, StructArray, UInt32Array},
+ buffer::Buffer,
+ datatypes::Field,
+ datatypes::ToByteSlice,
+ };
+
+ use super::*;
+
+ fn create_from_buffers() -> MapArray {
+ // Construct key and values
+ let keys_data = ArrayData::builder(DataType::Int32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+ let values_data = ArrayData::builder(DataType::UInt32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0u32, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
Review comment:
it might improve this test to use different values for the keys and values, so that any accidental mixup between them could be detected. Perhaps `[0, 10, 20, 30, 40, 50, 60, 70]`
##########
File path: arrow/src/array/array_map.rs
##########
@@ -0,0 +1,380 @@
+// 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 std::any::Any;
+use std::fmt;
+use std::mem;
+
+use super::make_array;
+use super::{
+ array::print_long_array, raw_pointer::RawPtrBox, Array, ArrayData, ArrayRef,
+};
+use crate::datatypes::{ArrowNativeType, DataType};
+use crate::error::ArrowError;
+
+/// A nested array type where each record is a key-value map.
+/// Keys should always be non-null, but values can be null.
+///
+/// [MapArray] is physically a [ListArray] that has a [StructArray]
+/// with 2 child fields.
+pub struct MapArray {
+ data: ArrayData,
+ values: ArrayRef,
+ value_offsets: RawPtrBox<i32>,
+}
+
+impl MapArray {
+ #[cfg(feature = "experimental-api")]
+ /// Returns a reference to the keys of this map.
+ pub fn keys(&self) -> ArrayRef {
+ make_array(self.values.data().child_data()[0].clone())
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns a reference to the values of this map.
+ pub fn values(&self) -> ArrayRef {
+ make_array(self.values.data().child_data()[1].clone())
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns the data type of the map's keys.
+ pub fn key_type(&self) -> DataType {
+ self.values.data().child_data()[0].data_type().clone()
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns the data type of the map's values.
+ pub fn value_type(&self) -> DataType {
+ self.values.data().child_data()[1].data_type().clone()
+ }
+
+ /// Returns ith value of this map array.
+ /// # Safety
+ /// Caller must ensure that the index is within the array bounds
+ pub unsafe fn value_unchecked(&self, i: usize) -> ArrayRef {
+ let end = *self.value_offsets().get_unchecked(i + 1);
+ let start = *self.value_offsets().get_unchecked(i);
+ self.values
+ .slice(start.to_usize().unwrap(), (end - start).to_usize().unwrap())
+ }
+
+ /// Returns ith value of this map array.
+ pub fn value(&self, i: usize) -> ArrayRef {
+ let end = self.value_offsets()[i + 1] as usize;
+ let start = self.value_offsets()[i] as usize;
+ self.values.slice(start, end - start)
+ }
+
+ /// Returns the offset values in the offsets buffer
+ #[inline]
+ pub fn value_offsets(&self) -> &[i32] {
+ // Soundness
+ // pointer alignment & location is ensured by RawPtrBox
+ // buffer bounds/offset is ensured by the ArrayData instance.
+ unsafe {
+ std::slice::from_raw_parts(
+ self.value_offsets.as_ptr().add(self.data.offset()),
+ self.len() + 1,
+ )
+ }
+ }
+
+ /// Returns the length for value at index `i`.
+ #[inline]
+ pub fn value_length(&self, i: usize) -> i32 {
+ let offsets = self.value_offsets();
+ offsets[i + 1] - offsets[i]
+ }
+}
+
+impl From<ArrayData> for MapArray {
+ fn from(data: ArrayData) -> Self {
+ Self::try_new_from_array_data(data)
+ .expect("Expected infallable creation of MapArray from ArrayData failed")
+ }
+}
+
+impl MapArray {
+ fn try_new_from_array_data(data: ArrayData) -> Result<Self, ArrowError> {
+ if data.buffers().len() != 1 {
+ return Err(ArrowError::InvalidArgumentError(
+ format!("MapArray data should contain a single buffer only (value offsets), had {}",
+ data.len())));
+ }
+
+ if data.child_data().len() != 1 {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a single child array (values array), had {}",
+ data.child_data().len()
+ )));
+ }
+
+ let entries = data.child_data()[0].clone();
+
+ if let DataType::Struct(fields) = entries.data_type() {
+ if fields.len() != 2 {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a struct array with 2 fields, have {} fields",
+ fields.len()
+ )));
+ }
+ } else {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a struct array child, found {:?}",
+ entries.data_type()
+ )));
+ }
+
+ let values = make_array(entries);
+ let value_offsets = data.buffers()[0].as_ptr();
+
+ let value_offsets = unsafe { RawPtrBox::<i32>::new(value_offsets) };
+ unsafe {
+ if (*value_offsets.as_ptr().offset(0)) != 0 {
+ return Err(ArrowError::InvalidArgumentError(String::from(
+ "offsets do not start at zero",
+ )));
+ }
+ }
+ Ok(Self {
+ data,
+ values,
+ value_offsets,
+ })
+ }
+}
+
+impl Array for MapArray {
+ fn as_any(&self) -> &Any {
+ self
+ }
+
+ fn data(&self) -> &ArrayData {
+ &self.data
+ }
+
+ /// Returns the total number of bytes of memory occupied by the buffers owned by this [MapArray].
+ fn get_buffer_memory_size(&self) -> usize {
+ self.data.get_buffer_memory_size()
+ }
+
+ /// Returns the total number of bytes of memory occupied physically by this [MapArray].
+ fn get_array_memory_size(&self) -> usize {
+ self.data.get_array_memory_size() + mem::size_of_val(self)
+ }
+}
+
+impl fmt::Debug for MapArray {
+ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(f, "MapArray\n[\n")?;
+ print_long_array(self, f, |array, index, f| {
+ fmt::Debug::fmt(&array.value(index), f)
+ })?;
+ write!(f, "]")
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use crate::{
+ array::ArrayData,
+ array::{Int32Array, StructArray, UInt32Array},
+ buffer::Buffer,
+ datatypes::Field,
+ datatypes::ToByteSlice,
+ };
+
+ use super::*;
+
+ fn create_from_buffers() -> MapArray {
+ // Construct key and values
+ let keys_data = ArrayData::builder(DataType::Int32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+ let values_data = ArrayData::builder(DataType::UInt32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0u32, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+
+ // Construct a buffer for value offsets, for the nested array:
+ // [[0, 1, 2], [3, 4, 5], [6, 7]]
+ let entry_offsets = Buffer::from(&[0, 3, 6, 8].to_byte_slice());
+
+ let keys = Field::new("keys", DataType::Int32, false);
+ let values = Field::new("values", DataType::UInt32, false);
+ let entry_struct = StructArray::from(vec![
+ (keys, make_array(keys_data)),
+ (values, make_array(values_data)),
+ ]);
+
+ // Construct a map array from the above two
+ let map_data_type = DataType::Map(
+ Box::new(Field::new(
+ "entries",
+ entry_struct.data_type().clone(),
+ true,
+ )),
+ false,
+ );
+ let map_data = ArrayData::builder(map_data_type)
+ .len(3)
+ .add_buffer(entry_offsets)
+ .add_child_data(entry_struct.data().clone())
+ .build();
+ MapArray::from(map_data)
+ }
+
+ #[test]
+ fn test_map_array() {
+ // Construct key and values
+ let key_data = ArrayData::builder(DataType::Int32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+ let value_data = ArrayData::builder(DataType::UInt32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0u32, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+
+ // Construct a buffer for value offsets, for the nested array:
+ // [[0, 1, 2], [3, 4, 5], [6, 7]]
+ let entry_offsets = Buffer::from(&[0, 3, 6, 8].to_byte_slice());
+
+ let keys_field = Field::new("keys", DataType::Int32, false);
+ let values_field = Field::new("values", DataType::UInt32, false);
+ let entry_struct = StructArray::from(vec![
+ (keys_field.clone(), make_array(key_data)),
+ (values_field.clone(), make_array(value_data.clone())),
+ ]);
+
+ // Construct a map array from the above two
+ let map_data_type = DataType::Map(
+ Box::new(Field::new(
+ "entries",
+ entry_struct.data_type().clone(),
+ true,
+ )),
+ false,
+ );
+ let map_data = ArrayData::builder(map_data_type)
+ .len(3)
+ .add_buffer(entry_offsets)
+ .add_child_data(entry_struct.data().clone())
+ .build();
+ let map_array = MapArray::from(map_data);
Review comment:
up to here, this method seems almost identical to `create_from_buffers` -- perhaps it could be reused?
##########
File path: parquet/src/arrow/arrow_reader.rs
##########
@@ -668,4 +668,20 @@ mod tests {
batch.unwrap();
}
}
+
+ #[test]
+ fn test_read_maps() {
+ let testdata = arrow::util::test_util::parquet_test_data();
+ let path = format!("{}/nested_maps.snappy.parquet", testdata);
Review comment:
nice!
##########
File path: parquet/src/arrow/levels.rs
##########
@@ -234,13 +234,53 @@ impl LevelInfo {
LevelType::Primitive(list_field.is_nullable()),
)]
}
- DataType::List(_) | DataType::LargeList(_) | DataType::Struct(_) => {
+ DataType::List(_)
+ | DataType::LargeList(_)
+ | DataType::Struct(_)
+ | DataType::Map(_, _) => {
list_level.calculate_array_levels(&child_array, list_field)
}
DataType::FixedSizeList(_, _) => unimplemented!(),
DataType::Union(_) => unimplemented!(),
}
}
+ DataType::Map(map_field, _) => {
+ // Calculate the map level
+ let map_level = self.calculate_child_levels(
+ array_offsets,
+ array_mask,
+ // A map is treated like a list as it has repetition
+ LevelType::List(field.is_nullable()),
+ );
+
+ let map_array = array.as_any().downcast_ref::<MapArray>().unwrap();
+
+ let key_array = map_array.keys();
+ let value_array = map_array.values();
+
+ if let DataType::Struct(fields) = map_field.data_type() {
+ let key_field = &fields[0];
+ let value_field = &fields[1];
+
+ let mut map_levels = vec![];
+
+ // Get key levels
+ let mut key_levels =
+ map_level.calculate_array_levels(&key_array, key_field);
+ map_levels.append(&mut key_levels);
+
+ let mut value_levels =
+ map_level.calculate_array_levels(&value_array, value_field);
+ map_levels.append(&mut value_levels);
+
+ map_levels
+ } else {
+ panic!(
+ "Map field should be a struct, found {:?}",
+ map_field.data_type()
+ );
+ }
+ }
DataType::FixedSizeList(_, _) => unimplemented!(),
Review comment:
I wonder if it would be useful to test in this module as well
##########
File path: arrow/src/json/reader.rs
##########
@@ -2177,6 +2279,81 @@ mod tests {
assert_eq!(read.data_ref(), expected.data_ref());
}
+ #[test]
+ fn test_map_json_arrays() {
+ let account_field = Field::new("account", DataType::UInt16, false);
+ let value_list_type =
+ DataType::List(Box::new(Field::new("item", DataType::Utf8, false)));
+ let entries_struct_type = DataType::Struct(vec![
+ Field::new("key", DataType::Utf8, false),
+ Field::new("value", value_list_type.clone(), true),
+ ]);
+ let stocks_field = Field::new(
+ "stocks",
+ DataType::Map(
+ Box::new(Field::new("entries", entries_struct_type.clone(), false)),
+ false,
+ ),
+ true,
+ );
+ let schema = Arc::new(Schema::new(vec![account_field, stocks_field.clone()]));
+ let builder = ReaderBuilder::new().with_schema(schema).with_batch_size(64);
+ // Note: account 456 has 'long' twice, to show that the JSON reader will overwrite
+ // existing keys. This thus guarantees unique keys for the map
Review comment:
👍
##########
File path: parquet/src/arrow/arrow_reader.rs
##########
@@ -668,4 +668,20 @@ mod tests {
batch.unwrap();
}
}
+
+ #[test]
+ fn test_read_maps() {
+ let testdata = arrow::util::test_util::parquet_test_data();
+ let path = format!("{}/nested_maps.snappy.parquet", testdata);
+ let parquet_file_reader =
+ SerializedFileReader::try_from(File::open(&path).unwrap()).unwrap();
+ let mut arrow_reader = ParquetFileArrowReader::new(Arc::new(parquet_file_reader));
+ let record_batch_reader = arrow_reader
+ .get_record_reader(60)
+ .expect("Failed to read into array!");
+
+ for batch in record_batch_reader {
+ batch.unwrap();
Review comment:
It seems like it would be valuable to check the contents of the record batches that were read out (in addition to checking they are not errors)
##########
File path: arrow/src/array/array.rs
##########
@@ -452,6 +453,7 @@ pub fn new_null_array(data_type: &DataType, length: usize) -> ArrayRef {
.map(|field| ArrayData::new_empty(field.data_type()))
.collect(),
)),
+ DataType::Map(_field, _keys_sorted) => todo!(),
Review comment:
should this still be "todo"?
##########
File path: arrow/src/array/array_map.rs
##########
@@ -0,0 +1,380 @@
+// 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 std::any::Any;
+use std::fmt;
+use std::mem;
+
+use super::make_array;
+use super::{
+ array::print_long_array, raw_pointer::RawPtrBox, Array, ArrayData, ArrayRef,
+};
+use crate::datatypes::{ArrowNativeType, DataType};
+use crate::error::ArrowError;
+
+/// A nested array type where each record is a key-value map.
+/// Keys should always be non-null, but values can be null.
+///
+/// [MapArray] is physically a [ListArray] that has a [StructArray]
+/// with 2 child fields.
+pub struct MapArray {
+ data: ArrayData,
+ values: ArrayRef,
+ value_offsets: RawPtrBox<i32>,
+}
+
+impl MapArray {
+ #[cfg(feature = "experimental-api")]
+ /// Returns a reference to the keys of this map.
+ pub fn keys(&self) -> ArrayRef {
+ make_array(self.values.data().child_data()[0].clone())
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns a reference to the values of this map.
+ pub fn values(&self) -> ArrayRef {
+ make_array(self.values.data().child_data()[1].clone())
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns the data type of the map's keys.
+ pub fn key_type(&self) -> DataType {
+ self.values.data().child_data()[0].data_type().clone()
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns the data type of the map's values.
+ pub fn value_type(&self) -> DataType {
+ self.values.data().child_data()[1].data_type().clone()
+ }
+
+ /// Returns ith value of this map array.
+ /// # Safety
+ /// Caller must ensure that the index is within the array bounds
+ pub unsafe fn value_unchecked(&self, i: usize) -> ArrayRef {
+ let end = *self.value_offsets().get_unchecked(i + 1);
+ let start = *self.value_offsets().get_unchecked(i);
+ self.values
+ .slice(start.to_usize().unwrap(), (end - start).to_usize().unwrap())
+ }
+
+ /// Returns ith value of this map array.
+ pub fn value(&self, i: usize) -> ArrayRef {
+ let end = self.value_offsets()[i + 1] as usize;
+ let start = self.value_offsets()[i] as usize;
+ self.values.slice(start, end - start)
+ }
+
+ /// Returns the offset values in the offsets buffer
+ #[inline]
+ pub fn value_offsets(&self) -> &[i32] {
+ // Soundness
+ // pointer alignment & location is ensured by RawPtrBox
+ // buffer bounds/offset is ensured by the ArrayData instance.
+ unsafe {
+ std::slice::from_raw_parts(
+ self.value_offsets.as_ptr().add(self.data.offset()),
+ self.len() + 1,
+ )
+ }
+ }
+
+ /// Returns the length for value at index `i`.
+ #[inline]
+ pub fn value_length(&self, i: usize) -> i32 {
+ let offsets = self.value_offsets();
+ offsets[i + 1] - offsets[i]
+ }
+}
+
+impl From<ArrayData> for MapArray {
+ fn from(data: ArrayData) -> Self {
+ Self::try_new_from_array_data(data)
+ .expect("Expected infallable creation of MapArray from ArrayData failed")
+ }
+}
+
+impl MapArray {
+ fn try_new_from_array_data(data: ArrayData) -> Result<Self, ArrowError> {
+ if data.buffers().len() != 1 {
+ return Err(ArrowError::InvalidArgumentError(
+ format!("MapArray data should contain a single buffer only (value offsets), had {}",
+ data.len())));
+ }
+
+ if data.child_data().len() != 1 {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a single child array (values array), had {}",
+ data.child_data().len()
+ )));
+ }
+
+ let entries = data.child_data()[0].clone();
+
+ if let DataType::Struct(fields) = entries.data_type() {
+ if fields.len() != 2 {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a struct array with 2 fields, have {} fields",
+ fields.len()
+ )));
+ }
+ } else {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a struct array child, found {:?}",
+ entries.data_type()
+ )));
+ }
+
+ let values = make_array(entries);
+ let value_offsets = data.buffers()[0].as_ptr();
+
+ let value_offsets = unsafe { RawPtrBox::<i32>::new(value_offsets) };
+ unsafe {
+ if (*value_offsets.as_ptr().offset(0)) != 0 {
+ return Err(ArrowError::InvalidArgumentError(String::from(
+ "offsets do not start at zero",
+ )));
+ }
+ }
+ Ok(Self {
+ data,
+ values,
+ value_offsets,
+ })
+ }
+}
+
+impl Array for MapArray {
+ fn as_any(&self) -> &Any {
+ self
+ }
+
+ fn data(&self) -> &ArrayData {
+ &self.data
+ }
+
+ /// Returns the total number of bytes of memory occupied by the buffers owned by this [MapArray].
+ fn get_buffer_memory_size(&self) -> usize {
+ self.data.get_buffer_memory_size()
+ }
+
+ /// Returns the total number of bytes of memory occupied physically by this [MapArray].
+ fn get_array_memory_size(&self) -> usize {
+ self.data.get_array_memory_size() + mem::size_of_val(self)
+ }
+}
+
+impl fmt::Debug for MapArray {
+ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(f, "MapArray\n[\n")?;
+ print_long_array(self, f, |array, index, f| {
+ fmt::Debug::fmt(&array.value(index), f)
+ })?;
+ write!(f, "]")
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use crate::{
+ array::ArrayData,
+ array::{Int32Array, StructArray, UInt32Array},
+ buffer::Buffer,
+ datatypes::Field,
+ datatypes::ToByteSlice,
+ };
+
+ use super::*;
+
+ fn create_from_buffers() -> MapArray {
+ // Construct key and values
+ let keys_data = ArrayData::builder(DataType::Int32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+ let values_data = ArrayData::builder(DataType::UInt32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0u32, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+
+ // Construct a buffer for value offsets, for the nested array:
+ // [[0, 1, 2], [3, 4, 5], [6, 7]]
+ let entry_offsets = Buffer::from(&[0, 3, 6, 8].to_byte_slice());
+
+ let keys = Field::new("keys", DataType::Int32, false);
+ let values = Field::new("values", DataType::UInt32, false);
+ let entry_struct = StructArray::from(vec![
+ (keys, make_array(keys_data)),
+ (values, make_array(values_data)),
+ ]);
+
+ // Construct a map array from the above two
+ let map_data_type = DataType::Map(
+ Box::new(Field::new(
+ "entries",
+ entry_struct.data_type().clone(),
+ true,
+ )),
+ false,
+ );
+ let map_data = ArrayData::builder(map_data_type)
+ .len(3)
+ .add_buffer(entry_offsets)
+ .add_child_data(entry_struct.data().clone())
+ .build();
+ MapArray::from(map_data)
+ }
+
+ #[test]
+ fn test_map_array() {
+ // Construct key and values
+ let key_data = ArrayData::builder(DataType::Int32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+ let value_data = ArrayData::builder(DataType::UInt32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0u32, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+
+ // Construct a buffer for value offsets, for the nested array:
+ // [[0, 1, 2], [3, 4, 5], [6, 7]]
+ let entry_offsets = Buffer::from(&[0, 3, 6, 8].to_byte_slice());
+
+ let keys_field = Field::new("keys", DataType::Int32, false);
+ let values_field = Field::new("values", DataType::UInt32, false);
+ let entry_struct = StructArray::from(vec![
+ (keys_field.clone(), make_array(key_data)),
+ (values_field.clone(), make_array(value_data.clone())),
+ ]);
+
+ // Construct a map array from the above two
+ let map_data_type = DataType::Map(
+ Box::new(Field::new(
+ "entries",
+ entry_struct.data_type().clone(),
+ true,
+ )),
+ false,
+ );
+ let map_data = ArrayData::builder(map_data_type)
+ .len(3)
+ .add_buffer(entry_offsets)
+ .add_child_data(entry_struct.data().clone())
+ .build();
+ let map_array = MapArray::from(map_data);
+
+ #[cfg(feature = "experimental-api")]
+ {
+ let values = map_array.values();
+ assert_eq!(&value_data, values.data());
+ assert_eq!(DataType::UInt32, map_array.value_type());
+ }
+ assert_eq!(3, map_array.len());
+ assert_eq!(0, map_array.null_count());
+ assert_eq!(6, map_array.value_offsets()[2]);
+ assert_eq!(2, map_array.value_length(2));
+
+ let key_array = Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef;
+ let value_array = Arc::new(UInt32Array::from(vec![0, 1, 2])) as ArrayRef;
+ let struct_array = StructArray::from(vec![
+ (keys_field.clone(), key_array),
+ (values_field.clone(), value_array),
+ ]);
+ assert_eq!(
+ struct_array,
+ StructArray::from(map_array.value(0).data().clone())
+ );
+ assert_eq!(
+ &struct_array,
+ unsafe { map_array.value_unchecked(0) }
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap()
+ );
+ for i in 0..3 {
+ assert!(map_array.is_valid(i));
+ assert!(!map_array.is_null(i));
+ }
+
+ // Now test with a non-zero offset
+ let map_data = ArrayData::builder(map_array.data_type().clone())
+ .len(3)
+ .offset(1)
+ .add_buffer(map_array.data().buffers()[0].clone())
+ .add_child_data(map_array.data().child_data()[0].clone())
+ .build();
+ let map_array = MapArray::from(map_data);
+
+ #[cfg(feature = "experimental-api")]
+ {
+ let values = map_array.values();
+ assert_eq!(&value_data, values.data());
+ assert_eq!(DataType::UInt32, map_array.value_type());
+ }
+ assert_eq!(3, map_array.len());
+ assert_eq!(0, map_array.null_count());
+ assert_eq!(6, map_array.value_offsets()[1]);
+ assert_eq!(2, map_array.value_length(1));
+
+ let key_array = Arc::new(Int32Array::from(vec![3, 4, 5])) as ArrayRef;
+ let value_array = Arc::new(UInt32Array::from(vec![3, 4, 5])) as ArrayRef;
+ let struct_array =
+ StructArray::from(vec![(keys_field, key_array), (values_field, value_array)]);
+ assert_eq!(
+ &struct_array,
+ map_array
+ .value(0)
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap()
+ );
+ assert_eq!(
+ &struct_array,
+ unsafe { map_array.value_unchecked(0) }
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap()
+ );
+ }
+
+ #[test]
+ fn test_map_array_slice() {
+ let map_array = create_from_buffers();
+
+ let sliced_array = map_array.slice(1, 2);
+ assert_eq!(2, sliced_array.len());
+ assert_eq!(1, sliced_array.offset());
+
+ // Check offset and length for each non-null value.
+ let sliced_map_array = sliced_array.as_any().downcast_ref::<MapArray>().unwrap();
+ assert_eq!(3, sliced_map_array.value_offsets()[0]);
+ assert_eq!(3, sliced_map_array.value_length(0));
+ assert_eq!(6, sliced_map_array.value_offsets()[1]);
+ assert_eq!(2, sliced_map_array.value_length(1));
+ }
Review comment:
it seems like it would also be good to check that the values of the slices are also correct.
##########
File path: arrow/src/array/transform/structure.rs
##########
@@ -26,13 +26,10 @@ pub(super) fn build_extend(array: &ArrayData) -> Extend {
index: usize,
start: usize,
len: usize| {
- mutable.child_data.iter_mut().for_each(|child| {
- child.extend(
- index,
- array.offset() + start,
Review comment:
why not use the `array.offset()`? I don't understand this change
##########
File path: arrow/src/array/array_map.rs
##########
@@ -0,0 +1,380 @@
+// 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 std::any::Any;
+use std::fmt;
+use std::mem;
+
+use super::make_array;
+use super::{
+ array::print_long_array, raw_pointer::RawPtrBox, Array, ArrayData, ArrayRef,
+};
+use crate::datatypes::{ArrowNativeType, DataType};
+use crate::error::ArrowError;
+
+/// A nested array type where each record is a key-value map.
+/// Keys should always be non-null, but values can be null.
+///
+/// [MapArray] is physically a [ListArray] that has a [StructArray]
+/// with 2 child fields.
+pub struct MapArray {
+ data: ArrayData,
+ values: ArrayRef,
+ value_offsets: RawPtrBox<i32>,
+}
+
+impl MapArray {
+ #[cfg(feature = "experimental-api")]
+ /// Returns a reference to the keys of this map.
+ pub fn keys(&self) -> ArrayRef {
+ make_array(self.values.data().child_data()[0].clone())
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns a reference to the values of this map.
+ pub fn values(&self) -> ArrayRef {
+ make_array(self.values.data().child_data()[1].clone())
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns the data type of the map's keys.
+ pub fn key_type(&self) -> DataType {
+ self.values.data().child_data()[0].data_type().clone()
+ }
+
+ #[cfg(feature = "experimental-api")]
+ /// Returns the data type of the map's values.
+ pub fn value_type(&self) -> DataType {
+ self.values.data().child_data()[1].data_type().clone()
+ }
+
+ /// Returns ith value of this map array.
+ /// # Safety
+ /// Caller must ensure that the index is within the array bounds
+ pub unsafe fn value_unchecked(&self, i: usize) -> ArrayRef {
+ let end = *self.value_offsets().get_unchecked(i + 1);
+ let start = *self.value_offsets().get_unchecked(i);
+ self.values
+ .slice(start.to_usize().unwrap(), (end - start).to_usize().unwrap())
+ }
+
+ /// Returns ith value of this map array.
+ pub fn value(&self, i: usize) -> ArrayRef {
+ let end = self.value_offsets()[i + 1] as usize;
+ let start = self.value_offsets()[i] as usize;
+ self.values.slice(start, end - start)
+ }
+
+ /// Returns the offset values in the offsets buffer
+ #[inline]
+ pub fn value_offsets(&self) -> &[i32] {
+ // Soundness
+ // pointer alignment & location is ensured by RawPtrBox
+ // buffer bounds/offset is ensured by the ArrayData instance.
+ unsafe {
+ std::slice::from_raw_parts(
+ self.value_offsets.as_ptr().add(self.data.offset()),
+ self.len() + 1,
+ )
+ }
+ }
+
+ /// Returns the length for value at index `i`.
+ #[inline]
+ pub fn value_length(&self, i: usize) -> i32 {
+ let offsets = self.value_offsets();
+ offsets[i + 1] - offsets[i]
+ }
+}
+
+impl From<ArrayData> for MapArray {
+ fn from(data: ArrayData) -> Self {
+ Self::try_new_from_array_data(data)
+ .expect("Expected infallable creation of MapArray from ArrayData failed")
+ }
+}
+
+impl MapArray {
+ fn try_new_from_array_data(data: ArrayData) -> Result<Self, ArrowError> {
+ if data.buffers().len() != 1 {
+ return Err(ArrowError::InvalidArgumentError(
+ format!("MapArray data should contain a single buffer only (value offsets), had {}",
+ data.len())));
+ }
+
+ if data.child_data().len() != 1 {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a single child array (values array), had {}",
+ data.child_data().len()
+ )));
+ }
+
+ let entries = data.child_data()[0].clone();
+
+ if let DataType::Struct(fields) = entries.data_type() {
+ if fields.len() != 2 {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a struct array with 2 fields, have {} fields",
+ fields.len()
+ )));
+ }
+ } else {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "MapArray should contain a struct array child, found {:?}",
+ entries.data_type()
+ )));
+ }
+
+ let values = make_array(entries);
+ let value_offsets = data.buffers()[0].as_ptr();
+
+ let value_offsets = unsafe { RawPtrBox::<i32>::new(value_offsets) };
+ unsafe {
+ if (*value_offsets.as_ptr().offset(0)) != 0 {
+ return Err(ArrowError::InvalidArgumentError(String::from(
+ "offsets do not start at zero",
+ )));
+ }
+ }
+ Ok(Self {
+ data,
+ values,
+ value_offsets,
+ })
+ }
+}
+
+impl Array for MapArray {
+ fn as_any(&self) -> &Any {
+ self
+ }
+
+ fn data(&self) -> &ArrayData {
+ &self.data
+ }
+
+ /// Returns the total number of bytes of memory occupied by the buffers owned by this [MapArray].
+ fn get_buffer_memory_size(&self) -> usize {
+ self.data.get_buffer_memory_size()
+ }
+
+ /// Returns the total number of bytes of memory occupied physically by this [MapArray].
+ fn get_array_memory_size(&self) -> usize {
+ self.data.get_array_memory_size() + mem::size_of_val(self)
+ }
+}
+
+impl fmt::Debug for MapArray {
+ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(f, "MapArray\n[\n")?;
+ print_long_array(self, f, |array, index, f| {
+ fmt::Debug::fmt(&array.value(index), f)
+ })?;
+ write!(f, "]")
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use crate::{
+ array::ArrayData,
+ array::{Int32Array, StructArray, UInt32Array},
+ buffer::Buffer,
+ datatypes::Field,
+ datatypes::ToByteSlice,
+ };
+
+ use super::*;
+
+ fn create_from_buffers() -> MapArray {
+ // Construct key and values
+ let keys_data = ArrayData::builder(DataType::Int32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+ let values_data = ArrayData::builder(DataType::UInt32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0u32, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+
+ // Construct a buffer for value offsets, for the nested array:
+ // [[0, 1, 2], [3, 4, 5], [6, 7]]
+ let entry_offsets = Buffer::from(&[0, 3, 6, 8].to_byte_slice());
+
+ let keys = Field::new("keys", DataType::Int32, false);
+ let values = Field::new("values", DataType::UInt32, false);
+ let entry_struct = StructArray::from(vec![
+ (keys, make_array(keys_data)),
+ (values, make_array(values_data)),
+ ]);
+
+ // Construct a map array from the above two
+ let map_data_type = DataType::Map(
+ Box::new(Field::new(
+ "entries",
+ entry_struct.data_type().clone(),
+ true,
+ )),
+ false,
+ );
+ let map_data = ArrayData::builder(map_data_type)
+ .len(3)
+ .add_buffer(entry_offsets)
+ .add_child_data(entry_struct.data().clone())
+ .build();
+ MapArray::from(map_data)
+ }
+
+ #[test]
+ fn test_map_array() {
+ // Construct key and values
+ let key_data = ArrayData::builder(DataType::Int32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+ let value_data = ArrayData::builder(DataType::UInt32)
+ .len(8)
+ .add_buffer(Buffer::from(&[0u32, 1, 2, 3, 4, 5, 6, 7].to_byte_slice()))
+ .build();
+
+ // Construct a buffer for value offsets, for the nested array:
+ // [[0, 1, 2], [3, 4, 5], [6, 7]]
+ let entry_offsets = Buffer::from(&[0, 3, 6, 8].to_byte_slice());
+
+ let keys_field = Field::new("keys", DataType::Int32, false);
+ let values_field = Field::new("values", DataType::UInt32, false);
+ let entry_struct = StructArray::from(vec![
+ (keys_field.clone(), make_array(key_data)),
+ (values_field.clone(), make_array(value_data.clone())),
+ ]);
+
+ // Construct a map array from the above two
+ let map_data_type = DataType::Map(
+ Box::new(Field::new(
+ "entries",
+ entry_struct.data_type().clone(),
+ true,
+ )),
+ false,
+ );
+ let map_data = ArrayData::builder(map_data_type)
+ .len(3)
+ .add_buffer(entry_offsets)
+ .add_child_data(entry_struct.data().clone())
+ .build();
+ let map_array = MapArray::from(map_data);
+
+ #[cfg(feature = "experimental-api")]
+ {
+ let values = map_array.values();
+ assert_eq!(&value_data, values.data());
+ assert_eq!(DataType::UInt32, map_array.value_type());
+ }
+ assert_eq!(3, map_array.len());
+ assert_eq!(0, map_array.null_count());
+ assert_eq!(6, map_array.value_offsets()[2]);
+ assert_eq!(2, map_array.value_length(2));
+
+ let key_array = Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef;
+ let value_array = Arc::new(UInt32Array::from(vec![0, 1, 2])) as ArrayRef;
Review comment:
it would really be good to have the keys and values contain different numbers, I think
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org