Skip to content

Commit

Permalink
Add ArrayAccessor trait (#1948)
Browse files Browse the repository at this point in the history
  • Loading branch information
tustvold committed Jul 22, 2022
1 parent 5e3facf commit b06996b
Show file tree
Hide file tree
Showing 8 changed files with 186 additions and 311 deletions.
73 changes: 73 additions & 0 deletions arrow/src/array/array.rs
Expand Up @@ -297,6 +297,79 @@ impl Array for ArrayRef {
}
}

impl<'a, T: Array> Array for &'a T {
fn as_any(&self) -> &dyn Any {
T::as_any(self)
}

fn data(&self) -> &ArrayData {
T::data(self)
}

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

fn data_ref(&self) -> &ArrayData {
T::data_ref(self)
}

fn data_type(&self) -> &DataType {
T::data_type(self)
}

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

fn len(&self) -> usize {
T::len(self)
}

fn is_empty(&self) -> bool {
T::is_empty(self)
}

fn offset(&self) -> usize {
T::offset(self)
}

fn is_null(&self, index: usize) -> bool {
T::is_null(self, index)
}

fn is_valid(&self, index: usize) -> bool {
T::is_valid(self, index)
}

fn null_count(&self) -> usize {
T::null_count(self)
}

fn get_buffer_memory_size(&self) -> usize {
T::get_buffer_memory_size(self)
}

fn get_array_memory_size(&self) -> usize {
T::get_array_memory_size(self)
}

fn to_raw(
&self,
) -> Result<(*const ffi::FFI_ArrowArray, *const ffi::FFI_ArrowSchema)> {
T::to_raw(self)
}
}

/// A generic trait for accessing the values of an [`Array`]
pub trait ArrayAccessor: Array {
type Item: Send + Sync;

fn value(&self, index: usize) -> Self::Item;

unsafe fn value_unchecked(&self, index: usize) -> Self::Item;
}

/// Constructs an array using the input `data`.
/// Returns a reference-counted `Array` instance.
pub fn make_array(data: ArrayData) -> ArrayRef {
Expand Down
15 changes: 15 additions & 0 deletions arrow/src/array/array_binary.rs
Expand Up @@ -23,6 +23,7 @@ use super::{
array::print_long_array, raw_pointer::RawPtrBox, Array, ArrayData,
FixedSizeListArray, GenericBinaryIter, GenericListArray, OffsetSizeTrait,
};
use crate::array::array::ArrayAccessor;
pub use crate::array::DecimalIter;
use crate::buffer::Buffer;
use crate::error::{ArrowError, Result};
Expand Down Expand Up @@ -245,6 +246,20 @@ impl<OffsetSize: OffsetSizeTrait> Array for GenericBinaryArray<OffsetSize> {
}
}

impl<'a, OffsetSize: OffsetSizeTrait> ArrayAccessor
for &'a GenericBinaryArray<OffsetSize>
{
type Item = &'a [u8];

fn value(&self, index: usize) -> Self::Item {
GenericBinaryArray::value(self, index)
}

unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
GenericBinaryArray::value_unchecked(self, index)
}
}

impl<OffsetSize: OffsetSizeTrait> From<ArrayData> for GenericBinaryArray<OffsetSize> {
fn from(data: ArrayData) -> Self {
assert_eq!(
Expand Down
13 changes: 13 additions & 0 deletions arrow/src/array/array_boolean.rs
Expand Up @@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.

use crate::array::array::ArrayAccessor;
use std::borrow::Borrow;
use std::convert::From;
use std::iter::{FromIterator, IntoIterator};
Expand Down Expand Up @@ -157,6 +158,18 @@ impl Array for BooleanArray {
}
}

impl<'a> ArrayAccessor for &'a BooleanArray {
type Item = bool;

fn value(&self, index: usize) -> Self::Item {
BooleanArray::value(self, index)
}

unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
BooleanArray::value_unchecked(self, index)
}
}

impl From<Vec<bool>> for BooleanArray {
fn from(data: Vec<bool>) -> Self {
let mut mut_buf = MutableBuffer::new_null(data.len());
Expand Down
27 changes: 26 additions & 1 deletion arrow/src/array/array_list.rs
Expand Up @@ -24,6 +24,7 @@ use super::{
array::print_long_array, make_array, raw_pointer::RawPtrBox, Array, ArrayData,
ArrayRef, BooleanBufferBuilder, GenericListArrayIter, PrimitiveArray,
};
use crate::array::array::ArrayAccessor;
use crate::{
buffer::MutableBuffer,
datatypes::{ArrowNativeType, ArrowPrimitiveType, DataType, Field},
Expand Down Expand Up @@ -245,7 +246,7 @@ impl<OffsetSize: OffsetSizeTrait> GenericListArray<OffsetSize> {
}
}

impl<OffsetSize: 'static + OffsetSizeTrait> Array for GenericListArray<OffsetSize> {
impl<OffsetSize: OffsetSizeTrait> Array for GenericListArray<OffsetSize> {
fn as_any(&self) -> &dyn Any {
self
}
Expand All @@ -259,6 +260,18 @@ impl<OffsetSize: 'static + OffsetSizeTrait> Array for GenericListArray<OffsetSiz
}
}

impl<'a, OffsetSize: OffsetSizeTrait> ArrayAccessor for &'a GenericListArray<OffsetSize> {
type Item = ArrayRef;

fn value(&self, index: usize) -> Self::Item {
GenericListArray::value(self, index)
}

unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
GenericListArray::value(self, index)
}
}

impl<OffsetSize: OffsetSizeTrait> fmt::Debug for GenericListArray<OffsetSize> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let prefix = if OffsetSize::IS_LARGE { "Large" } else { "" };
Expand Down Expand Up @@ -466,6 +479,18 @@ impl Array for FixedSizeListArray {
}
}

impl ArrayAccessor for FixedSizeListArray {
type Item = ArrayRef;

fn value(&self, index: usize) -> Self::Item {
FixedSizeListArray::value(self, index)
}

unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
FixedSizeListArray::value(self, index)
}
}

impl fmt::Debug for FixedSizeListArray {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "FixedSizeListArray<{}>\n[\n", self.value_length())?;
Expand Down
13 changes: 13 additions & 0 deletions arrow/src/array/array_primitive.rs
Expand Up @@ -33,6 +33,7 @@ use crate::{
util::trusted_len_unzip,
};

use crate::array::array::ArrayAccessor;
use half::f16;

/// Array whose elements are of primitive types.
Expand Down Expand Up @@ -188,6 +189,18 @@ impl<T: ArrowPrimitiveType> Array for PrimitiveArray<T> {
}
}

impl<'a, T: ArrowPrimitiveType> ArrayAccessor for &'a PrimitiveArray<T> {
type Item = T::Native;

fn value(&self, index: usize) -> Self::Item {
PrimitiveArray::value(self, index)
}

unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
PrimitiveArray::value_unchecked(self, index)
}
}

fn as_datetime<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveDateTime> {
match T::DATA_TYPE {
DataType::Date32 => Some(temporal_conversions::date32_to_datetime(v as i32)),
Expand Down
15 changes: 15 additions & 0 deletions arrow/src/array/array_string.rs
Expand Up @@ -23,6 +23,7 @@ use super::{
array::print_long_array, raw_pointer::RawPtrBox, Array, ArrayData, GenericListArray,
GenericStringIter, OffsetSizeTrait,
};
use crate::array::array::ArrayAccessor;
use crate::buffer::Buffer;
use crate::util::bit_util;
use crate::{buffer::MutableBuffer, datatypes::DataType};
Expand Down Expand Up @@ -298,6 +299,20 @@ impl<OffsetSize: OffsetSizeTrait> Array for GenericStringArray<OffsetSize> {
}
}

impl<'a, OffsetSize: OffsetSizeTrait> ArrayAccessor
for &'a GenericStringArray<OffsetSize>
{
type Item = &'a str;

fn value(&self, index: usize) -> Self::Item {
GenericStringArray::value(self, index)
}

unsafe fn value_unchecked(&self, index: usize) -> Self::Item {
GenericStringArray::value_unchecked(self, index)
}
}

impl<OffsetSize: OffsetSizeTrait> From<ArrayData> for GenericStringArray<OffsetSize> {
fn from(data: ArrayData) -> Self {
assert_eq!(
Expand Down
10 changes: 10 additions & 0 deletions arrow/src/array/equal_json.rs
Expand Up @@ -37,6 +37,16 @@ pub trait JsonEqual {
}
}

impl<'a, T: JsonEqual> JsonEqual for &'a T {
fn equals_json(&self, json: &[&Value]) -> bool {
T::equals_json(self, json)
}

fn equals_json_values(&self, json: &[Value]) -> bool {
T::equals_json_values(self, json)
}
}

/// Implement array equals for numeric type
impl<T: ArrowPrimitiveType> JsonEqual for PrimitiveArray<T> {
fn equals_json(&self, json: &[&Value]) -> bool {
Expand Down

0 comments on commit b06996b

Please sign in to comment.