Skip to content
Permalink
Browse files
Allow primitive array creation from iterators of PrimitiveTypes (as w…
…ell as `Option`) (#1367)

* More idiomatic primitive array creation

* Use From instead for clippy

* Rename to NativeAdapter and add document
  • Loading branch information
viirya committed Mar 3, 2022
1 parent f9e4cf5 commit 258f828f9fd5bdfb668cbab7005d6bff73184b8f
Showing 1 changed file with 68 additions and 5 deletions.
@@ -16,7 +16,6 @@
// under the License.

use std::any::Any;
use std::borrow::Borrow;
use std::convert::From;
use std::fmt;
use std::iter::{FromIterator, IntoIterator};
@@ -34,6 +33,8 @@ use crate::{
util::trusted_len_unzip,
};

use half::f16;

/// Array whose elements are of primitive types.
///
/// # Example: From an iterator of values
@@ -340,8 +341,57 @@ impl<'a, T: ArrowPrimitiveType> PrimitiveArray<T> {
}
}

impl<T: ArrowPrimitiveType, Ptr: Borrow<Option<<T as ArrowPrimitiveType>::Native>>>
FromIterator<Ptr> for PrimitiveArray<T>
/// This struct is used as an adapter when creating `PrimitiveArray` from an iterator.
/// `FromIterator` for `PrimitiveArray` takes an iterator where the elements can be `into`
/// this struct. So once implementing `From` or `Into` trait for a type, an iterator of
/// the type can be collected to `PrimitiveArray`.
#[derive(Debug)]
pub struct NativeAdapter<T: ArrowPrimitiveType> {
pub native: Option<T::Native>,
}

macro_rules! def_from_for_primitive {
( $ty:ident, $tt:tt) => {
impl From<$tt> for NativeAdapter<$ty> {
fn from(value: $tt) -> Self {
NativeAdapter {
native: Some(value),
}
}
}
};
}

def_from_for_primitive!(Int8Type, i8);
def_from_for_primitive!(Int16Type, i16);
def_from_for_primitive!(Int32Type, i32);
def_from_for_primitive!(Int64Type, i64);
def_from_for_primitive!(UInt8Type, u8);
def_from_for_primitive!(UInt16Type, u16);
def_from_for_primitive!(UInt32Type, u32);
def_from_for_primitive!(UInt64Type, u64);
def_from_for_primitive!(Float16Type, f16);
def_from_for_primitive!(Float32Type, f32);
def_from_for_primitive!(Float64Type, f64);

impl<T: ArrowPrimitiveType> From<Option<<T as ArrowPrimitiveType>::Native>>
for NativeAdapter<T>
{
fn from(value: Option<<T as ArrowPrimitiveType>::Native>) -> Self {
NativeAdapter { native: value }
}
}

impl<T: ArrowPrimitiveType> From<&Option<<T as ArrowPrimitiveType>::Native>>
for NativeAdapter<T>
{
fn from(value: &Option<<T as ArrowPrimitiveType>::Native>) -> Self {
NativeAdapter { native: *value }
}
}

impl<'a, T: ArrowPrimitiveType, Ptr: Into<NativeAdapter<T>>> FromIterator<Ptr>
for PrimitiveArray<T>
{
fn from_iter<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
let iter = iter.into_iter();
@@ -351,9 +401,9 @@ impl<T: ArrowPrimitiveType, Ptr: Borrow<Option<<T as ArrowPrimitiveType>::Native

let buffer: Buffer = iter
.map(|item| {
if let Some(a) = item.borrow() {
if let Some(a) = item.into().native {
null_buf.append(true);
*a
a
} else {
null_buf.append(false);
// this ensures that null items on the buffer are not arbitrary.
@@ -530,6 +580,7 @@ mod tests {
use std::thread;

use crate::buffer::Buffer;
use crate::compute::eq_dyn;
use crate::datatypes::DataType;

#[test]
@@ -998,4 +1049,16 @@ mod tests {
assert!(ret.is_ok());
assert_eq!(8, ret.ok().unwrap());
}

#[test]
fn test_primitive_array_creation() {
let array1: Int8Array = [10_i8, 11, 12, 13, 14].into_iter().collect();
let array2: Int8Array = [10_i8, 11, 12, 13, 14].into_iter().map(Some).collect();

let result = eq_dyn(&array1, &array2);
assert_eq!(
result.unwrap(),
BooleanArray::from(vec![true, true, true, true, true])
);
}
}

0 comments on commit 258f828

Please sign in to comment.