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

Add Array::to_data and Array::nulls (#3880) #3881

Merged
merged 5 commits into from
Mar 17, 2023
Merged
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
8 changes: 4 additions & 4 deletions arrow-arith/src/aggregate.rs
Original file line number Diff line number Diff line change
@@ -117,7 +117,7 @@ where
.map(|i| unsafe { array.value_unchecked(i) })
.reduce(|acc, item| if cmp(&acc, &item) { item } else { acc })
} else {
let nulls = array.data().nulls().unwrap();
let nulls = array.nulls().unwrap();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The point of this PR, as I understand it, is to remove direct use of data() -- so that eventually the Arrays themselves do not need to hold ArrayData but rather the strongly typed variants.

The plan is not to remove ArrayData completely, but rather construct it on demand if needed from the appropriate parts of each Array

let iter = BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len());
unsafe {
let idx = iter.reduce(|acc_idx, idx| {
@@ -288,7 +288,7 @@ where

let data: &[T::Native] = array.values();

match array.data().nulls() {
match array.nulls() {
None => {
let sum = data.iter().fold(T::default_value(), |accumulator, value| {
accumulator.add_wrapping(*value)
@@ -347,7 +347,7 @@ where

let data: &[T::Native] = array.values();

match array.data().nulls() {
match array.nulls() {
None => {
let sum = data
.iter()
@@ -665,7 +665,7 @@ mod simd {
let mut chunk_acc = A::init_accumulator_chunk();
let mut rem_acc = A::init_accumulator_scalar();

match array.data().nulls() {
match array.nulls() {
None => {
let data_chunks = data.chunks_exact(64);
let remainder = data_chunks.remainder();
3 changes: 1 addition & 2 deletions arrow-arith/src/arity.rs
Original file line number Diff line number Diff line change
@@ -414,8 +414,7 @@ where

let array_builder = builder
.finish()
.data()
.clone()
.into_data()
.into_builder()
.null_bit_buffer(null_buffer)
.null_count(null_count);
24 changes: 12 additions & 12 deletions arrow-arith/src/boolean.rs
Original file line number Diff line number Diff line change
@@ -610,7 +610,7 @@ mod tests {
let a = BooleanArray::from(vec![false, false, false, true, true, true]);

// ensure null bitmap of a is absent
assert!(a.data().nulls().is_none());
assert!(a.nulls().is_none());

let b = BooleanArray::from(vec![
Some(true),
@@ -622,7 +622,7 @@ mod tests {
]);

// ensure null bitmap of b is present
assert!(b.data().nulls().is_some());
assert!(b.nulls().is_some());

let c = or_kleene(&a, &b).unwrap();

@@ -650,12 +650,12 @@ mod tests {
]);

// ensure null bitmap of b is absent
assert!(a.data().nulls().is_some());
assert!(a.nulls().is_some());

let b = BooleanArray::from(vec![false, false, false, true, true, true]);

// ensure null bitmap of a is present
assert!(b.data().nulls().is_none());
assert!(b.nulls().is_none());

let c = or_kleene(&a, &b).unwrap();

@@ -852,7 +852,7 @@ mod tests {
let expected = BooleanArray::from(vec![false, false, false, false]);

assert_eq!(expected, res);
assert!(res.data().nulls().is_none());
assert!(res.nulls().is_none());
}

#[test]
@@ -865,7 +865,7 @@ mod tests {
let expected = BooleanArray::from(vec![false, false, false, false]);

assert_eq!(expected, res);
assert!(res.data().nulls().is_none());
assert!(res.nulls().is_none());
}

#[test]
@@ -877,7 +877,7 @@ mod tests {
let expected = BooleanArray::from(vec![true, true, true, true]);

assert_eq!(expected, res);
assert!(res.data().nulls().is_none());
assert!(res.nulls().is_none());
}

#[test]
@@ -890,7 +890,7 @@ mod tests {
let expected = BooleanArray::from(vec![true, true, true, true]);

assert_eq!(expected, res);
assert!(res.data().nulls().is_none());
assert!(res.nulls().is_none());
}

#[test]
@@ -902,7 +902,7 @@ mod tests {
let expected = BooleanArray::from(vec![false, true, false, true]);

assert_eq!(expected, res);
assert!(res.data().nulls().is_none());
assert!(res.nulls().is_none());
}

#[test]
@@ -933,7 +933,7 @@ mod tests {
let expected = BooleanArray::from(vec![false, true, false, true]);

assert_eq!(expected, res);
assert!(res.data().nulls().is_none());
assert!(res.nulls().is_none());
}

#[test]
@@ -945,7 +945,7 @@ mod tests {
let expected = BooleanArray::from(vec![true, false, true, false]);

assert_eq!(expected, res);
assert!(res.data().nulls().is_none());
assert!(res.nulls().is_none());
}

#[test]
@@ -976,6 +976,6 @@ mod tests {
let expected = BooleanArray::from(vec![true, false, true, false]);

assert_eq!(expected, res);
assert!(res.data().nulls().is_none());
assert!(res.nulls().is_none());
}
}
2 changes: 1 addition & 1 deletion arrow-array/src/array/binary_array.rs
Original file line number Diff line number Diff line change
@@ -77,7 +77,7 @@ impl<OffsetSize: OffsetSizeTrait> GenericBinaryArray<OffsetSize> {
.offset(v.offset())
.add_buffer(v.data_ref().buffers()[0].clone())
.add_buffer(child_data.buffers()[0].slice(child_data.offset()))
.nulls(v.data().nulls().cloned());
.nulls(v.nulls().cloned());

let data = unsafe { builder.build_unchecked() };
Self::from(data)
17 changes: 16 additions & 1 deletion arrow-array/src/array/boolean_array.rs
Original file line number Diff line number Diff line change
@@ -18,12 +18,14 @@
use crate::array::print_long_array;
use crate::builder::BooleanBuilder;
use crate::iterator::BooleanIter;
use crate::{Array, ArrayAccessor};
use crate::{Array, ArrayAccessor, ArrayRef};
use arrow_buffer::buffer::NullBuffer;
use arrow_buffer::{bit_util, Buffer, MutableBuffer};
use arrow_data::bit_mask::combine_option_bitmap;
use arrow_data::ArrayData;
use arrow_schema::DataType;
use std::any::Any;
use std::sync::Arc;

/// Array of bools
///
@@ -265,9 +267,22 @@ impl Array for BooleanArray {
&self.data
}

fn to_data(&self) -> ArrayData {
self.data.clone()
}

fn into_data(self) -> ArrayData {
self.into()
}

fn slice(&self, offset: usize, length: usize) -> ArrayRef {
// TODO: Slice buffers directly (#3880)
Arc::new(Self::from(self.data.slice(offset, length)))
}

fn nulls(&self) -> Option<&NullBuffer> {
self.data.nulls()
}
}

impl<'a> ArrayAccessor for &'a BooleanArray {
18 changes: 16 additions & 2 deletions arrow-array/src/array/byte_array.rs
Original file line number Diff line number Diff line change
@@ -20,12 +20,13 @@ use crate::builder::GenericByteBuilder;
use crate::iterator::ArrayIter;
use crate::types::bytes::ByteArrayNativeType;
use crate::types::ByteArrayType;
use crate::{Array, ArrayAccessor, OffsetSizeTrait};
use arrow_buffer::OffsetBuffer;
use crate::{Array, ArrayAccessor, ArrayRef, OffsetSizeTrait};
use arrow_buffer::{ArrowNativeType, Buffer};
use arrow_buffer::{NullBuffer, OffsetBuffer};
use arrow_data::ArrayData;
use arrow_schema::DataType;
use std::any::Any;
use std::sync::Arc;

/// Generic struct for variable-size byte arrays
///
@@ -237,9 +238,22 @@ impl<T: ByteArrayType> Array for GenericByteArray<T> {
&self.data
}

fn to_data(&self) -> ArrayData {
self.data.clone()
}

fn into_data(self) -> ArrayData {
self.into()
}

fn slice(&self, offset: usize, length: usize) -> ArrayRef {
// TODO: Slice buffers directly (#3880)
Arc::new(Self::from(self.data.slice(offset, length)))
}

fn nulls(&self) -> Option<&NullBuffer> {
self.data.nulls()
}
}

impl<'a, T: ByteArrayType> ArrayAccessor for &'a GenericByteArray<T> {
27 changes: 27 additions & 0 deletions arrow-array/src/array/dictionary_array.rs
Original file line number Diff line number Diff line change
@@ -23,10 +23,12 @@ use crate::{
make_array, Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, PrimitiveArray,
StringArray,
};
use arrow_buffer::buffer::NullBuffer;
use arrow_buffer::ArrowNativeType;
use arrow_data::ArrayData;
use arrow_schema::{ArrowError, DataType};
use std::any::Any;
use std::sync::Arc;

///
/// A dictionary array where each element is a single value indexed by an integer key.
@@ -590,9 +592,22 @@ impl<T: ArrowDictionaryKeyType> Array for DictionaryArray<T> {
&self.data
}

fn to_data(&self) -> ArrayData {
self.data.clone()
}

fn into_data(self) -> ArrayData {
self.into()
}

fn slice(&self, offset: usize, length: usize) -> ArrayRef {
// TODO: Slice buffers directly (#3880)
Arc::new(Self::from(self.data.slice(offset, length)))
}

fn nulls(&self) -> Option<&NullBuffer> {
self.data.nulls()
}
}

impl<T: ArrowDictionaryKeyType> std::fmt::Debug for DictionaryArray<T> {
@@ -669,9 +684,21 @@ impl<'a, K: ArrowDictionaryKeyType, V: Sync> Array for TypedDictionaryArray<'a,
&self.dictionary.data
}

fn to_data(&self) -> ArrayData {
self.dictionary.to_data()
}

fn into_data(self) -> ArrayData {
self.dictionary.into_data()
}

fn slice(&self, offset: usize, length: usize) -> ArrayRef {
self.dictionary.slice(offset, length)
}

fn nulls(&self) -> Option<&NullBuffer> {
self.dictionary.nulls()
}
}

impl<'a, K, V> IntoIterator for TypedDictionaryArray<'a, K, V>
17 changes: 16 additions & 1 deletion arrow-array/src/array/fixed_size_binary_array.rs
Original file line number Diff line number Diff line change
@@ -17,11 +17,13 @@

use crate::array::print_long_array;
use crate::iterator::FixedSizeBinaryIter;
use crate::{Array, ArrayAccessor, FixedSizeListArray};
use crate::{Array, ArrayAccessor, ArrayRef, FixedSizeListArray};
use arrow_buffer::buffer::NullBuffer;
use arrow_buffer::{bit_util, Buffer, MutableBuffer};
use arrow_data::ArrayData;
use arrow_schema::{ArrowError, DataType};
use std::any::Any;
use std::sync::Arc;

/// An array where each element is a fixed-size sequence of bytes.
///
@@ -462,9 +464,22 @@ impl Array for FixedSizeBinaryArray {
&self.data
}

fn to_data(&self) -> ArrayData {
self.data.clone()
}

fn into_data(self) -> ArrayData {
self.into()
}

fn slice(&self, offset: usize, length: usize) -> ArrayRef {
// TODO: Slice buffers directly (#3880)
Arc::new(Self::from(self.data.slice(offset, length)))
}

fn nulls(&self) -> Option<&NullBuffer> {
self.data.nulls()
}
}

impl<'a> ArrayAccessor for &'a FixedSizeBinaryArray {
15 changes: 15 additions & 0 deletions arrow-array/src/array/fixed_size_list_array.rs
Original file line number Diff line number Diff line change
@@ -18,9 +18,11 @@
use crate::array::print_long_array;
use crate::builder::{FixedSizeListBuilder, PrimitiveBuilder};
use crate::{make_array, Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType};
use arrow_buffer::buffer::NullBuffer;
use arrow_data::ArrayData;
use arrow_schema::DataType;
use std::any::Any;
use std::sync::Arc;

/// A list array where each element is a fixed-size sequence of values with the same
/// type whose maximum length is represented by a i32.
@@ -205,9 +207,22 @@ impl Array for FixedSizeListArray {
&self.data
}

fn to_data(&self) -> ArrayData {
self.data.clone()
}

fn into_data(self) -> ArrayData {
self.into()
}

fn slice(&self, offset: usize, length: usize) -> ArrayRef {
// TODO: Slice buffers directly (#3880)
Arc::new(Self::from(self.data.slice(offset, length)))
}

fn nulls(&self) -> Option<&NullBuffer> {
self.data.nulls()
}
}

impl ArrayAccessor for FixedSizeListArray {
Loading