Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Return references from FixedSizeListArray and MapArray #3652

Merged
merged 3 commits into from
Feb 2, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 5 additions & 8 deletions arrow-array/src/array/fixed_size_list_array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,8 +69,8 @@ pub struct FixedSizeListArray {

impl FixedSizeListArray {
/// Returns a reference to the values of this list.
pub fn values(&self) -> ArrayRef {
self.values.clone()
pub fn values(&self) -> &ArrayRef {
&self.values
}

/// Returns a clone of the value type of this list.
Expand Down Expand Up @@ -261,8 +261,7 @@ mod tests {
.unwrap();
let list_array = FixedSizeListArray::from(list_data);

let values = list_array.values();
assert_eq!(&value_data, values.data());
assert_eq!(&value_data, list_array.values().data());
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(3, list_array.len());
assert_eq!(0, list_array.null_count());
Expand Down Expand Up @@ -291,8 +290,7 @@ mod tests {
.unwrap();
let list_array = FixedSizeListArray::from(list_data);

let values = list_array.values();
assert_eq!(&value_data, values.data());
assert_eq!(&value_data, list_array.values().data());
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(3, list_array.len());
assert_eq!(0, list_array.null_count());
Expand Down Expand Up @@ -368,8 +366,7 @@ mod tests {
.unwrap();
let list_array = FixedSizeListArray::from(list_data);

let values = list_array.values();
assert_eq!(&value_data, values.data());
assert_eq!(&value_data, list_array.values().data());
assert_eq!(DataType::Int32, list_array.value_type());
assert_eq!(5, list_array.len());
assert_eq!(2, list_array.null_count());
Expand Down
56 changes: 35 additions & 21 deletions arrow-array/src/array/map_array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,50 +27,60 @@ use std::sync::Arc;
/// Keys should always be non-null, but values can be null.
///
/// [MapArray] is physically a [crate::array::ListArray] that has a
/// [crate::array::StructArray] with 2 child fields.
/// [StructArray] with 2 child fields.
#[derive(Clone)]
pub struct MapArray {
data: ArrayData,
/// The [`StructArray`] that is the direct child of this array
entries: ArrayRef,
/// The first child of `entries`, the "keys" of this MapArray
keys: ArrayRef,
/// The second child of `entries`, the "values" of this MapArray
values: ArrayRef,
/// The start and end offsets of each entry
value_offsets: RawPtrBox<i32>,
}

impl MapArray {
/// Returns a reference to the keys of this map.
pub fn keys(&self) -> ArrayRef {
make_array(self.values.data().child_data()[0].clone())
pub fn keys(&self) -> &ArrayRef {
&self.keys
}

/// Returns a reference to the values of this map.
pub fn values(&self) -> ArrayRef {
make_array(self.values.data().child_data()[1].clone())
pub fn values(&self) -> &ArrayRef {
&self.values
}

/// Returns the data type of the map's keys.
pub fn key_type(&self) -> DataType {
self.values.data().child_data()[0].data_type().clone()
pub fn key_type(&self) -> &DataType {
self.keys.data_type()
}

/// Returns the data type of the map's values.
pub fn value_type(&self) -> DataType {
self.values.data().child_data()[1].data_type().clone()
pub fn value_type(&self) -> &DataType {
self.values.data_type()
}

/// Returns ith value of this map array.
///
/// This is a [`StructArray`] containing two fields
/// # 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
self.entries
.slice(start.to_usize().unwrap(), (end - start).to_usize().unwrap())
}

/// Returns ith value of this map array.
///
/// This is a [`StructArray`] containing two fields
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)
self.entries.slice(start, end - start)
}

/// Returns the offset values in the offsets buffer
Expand Down Expand Up @@ -146,7 +156,9 @@ impl MapArray {
)));
}

let values = make_array(entries);
let keys = make_array(entries.child_data()[0].clone());
let values = make_array(entries.child_data()[1].clone());
let entries = make_array(entries);
let value_offsets = data.buffers()[0].as_ptr();

// SAFETY:
Expand All @@ -159,8 +171,11 @@ impl MapArray {
)));
}
}

Ok(Self {
data,
entries,
keys,
values,
value_offsets,
})
Expand Down Expand Up @@ -241,6 +256,8 @@ impl std::fmt::Debug for MapArray {

#[cfg(test)]
mod tests {
use crate::cast::as_primitive_array;
use crate::types::UInt32Type;
use crate::{Int32Array, UInt32Array};
use std::sync::Arc;

Expand Down Expand Up @@ -335,9 +352,8 @@ mod tests {
.unwrap();
let map_array = MapArray::from(map_data);

let values = map_array.values();
assert_eq!(&value_data, values.data());
assert_eq!(DataType::UInt32, map_array.value_type());
assert_eq!(&value_data, map_array.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]);
Expand Down Expand Up @@ -376,9 +392,8 @@ mod tests {
.unwrap();
let map_array = MapArray::from(map_data);

let values = map_array.values();
assert_eq!(&value_data, values.data());
assert_eq!(DataType::UInt32, map_array.value_type());
assert_eq!(&value_data, map_array.values().data());
assert_eq!(&DataType::UInt32, map_array.value_type());
assert_eq!(2, map_array.len());
assert_eq!(0, map_array.null_count());
assert_eq!(6, map_array.value_offsets()[1]);
Expand Down Expand Up @@ -508,12 +523,11 @@ mod tests {
)
.unwrap();

let values = map_array.values();
assert_eq!(
&values_data,
values.as_any().downcast_ref::<UInt32Array>().unwrap()
as_primitive_array::<UInt32Type>(map_array.values())
);
assert_eq!(DataType::UInt32, map_array.value_type());
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]);
Expand Down
4 changes: 2 additions & 2 deletions arrow-array/src/builder/map_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,8 @@ use std::sync::Arc;
///
/// let array = builder.finish();
/// assert_eq!(array.value_offsets(), &[0, 1, 3, 3, 3]);
/// assert_eq!(*array.values(), Int32Array::from(vec![1, 2, 4]));
/// assert_eq!(*array.keys(), StringArray::from(vec!["joe", "blogs", "foo"]));
/// assert_eq!(array.values().as_ref(), &Int32Array::from(vec![1, 2, 4]));
/// assert_eq!(array.keys().as_ref(), &StringArray::from(vec!["joe", "blogs", "foo"]));
///
/// ```
#[derive(Debug)]
Expand Down
6 changes: 3 additions & 3 deletions arrow-ipc/src/writer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,7 @@ impl IpcDataGenerator {
.expect("Unable to downcast to fixed size list array");
self.encode_dictionaries(
field,
&list.values(),
list.values(),
encoded_dictionaries,
dictionary_tracker,
write_options,
Expand All @@ -264,7 +264,7 @@ impl IpcDataGenerator {
// keys
self.encode_dictionaries(
keys,
&map_array.keys(),
map_array.keys(),
encoded_dictionaries,
dictionary_tracker,
write_options,
Expand All @@ -273,7 +273,7 @@ impl IpcDataGenerator {
// values
self.encode_dictionaries(
values,
&map_array.values(),
map_array.values(),
encoded_dictionaries,
dictionary_tracker,
write_options,
Expand Down
4 changes: 2 additions & 2 deletions arrow-json/src/writer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -415,8 +415,8 @@ fn set_column_for_json_rows(
)));
}

let keys = as_string_array(&keys);
let values = array_to_json_array(&values)?;
let keys = as_string_array(keys);
let values = array_to_json_array(values)?;

let mut kv = keys.iter().zip(values.into_iter());

Expand Down
4 changes: 2 additions & 2 deletions parquet/src/arrow/arrow_writer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -333,8 +333,8 @@ fn write_leaves<W: Write>(
.as_any()
.downcast_ref::<arrow_array::MapArray>()
.expect("Unable to get map array");
keys.push(map_array.keys());
values.push(map_array.values());
keys.push(map_array.keys().clone());
values.push(map_array.values().clone());
}

write_leaves(row_group_writer, &keys, levels)?;
Expand Down