From d278f080bc7e74879d103d393d755e567c3b68c6 Mon Sep 17 00:00:00 2001 From: Jorge Leitao Date: Sun, 22 Aug 2021 20:42:09 +0100 Subject: [PATCH] Added trait TryPush. (#314) --- src/array/binary/mutable.rs | 63 ++++++++++++++++++---------------- src/array/boolean/mutable.rs | 11 ++++-- src/array/list/mutable.rs | 28 ++++++++++----- src/array/mod.rs | 6 ++++ src/array/primitive/mutable.rs | 10 +++++- src/array/utf8/mutable.rs | 59 ++++++++++++++++--------------- tests/it/array/list/mutable.rs | 8 +++++ 7 files changed, 115 insertions(+), 70 deletions(-) diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs index 03264972781..fe4b3e050fa 100644 --- a/src/array/binary/mutable.rs +++ b/src/array/binary/mutable.rs @@ -1,7 +1,7 @@ use std::{iter::FromIterator, sync::Arc}; use crate::{ - array::{Array, MutableArray, Offset, TryExtend}, + array::{Array, MutableArray, Offset, TryExtend, TryPush}, bitmap::MutableBitmap, buffer::MutableBuffer, datatypes::DataType, @@ -61,34 +61,9 @@ impl MutableBinaryArray { *self.offsets.last().unwrap() } - pub fn try_push>(&mut self, value: Option) -> Result<()> { - match value { - Some(value) => { - let bytes = value.as_ref(); - - let size = O::from_usize(self.values.len() + bytes.len()) - .ok_or(ArrowError::KeyOverflowError)?; - - self.values.extend_from_slice(bytes); - - self.offsets.push(size); - - match &mut self.validity { - Some(validity) => validity.push(true), - None => {} - } - } - None => { - self.offsets.push(self.last_offset()); - match &mut self.validity { - Some(validity) => validity.push(false), - None => self.init_validity(), - } - } - } - Ok(()) - } - + /// Pushes a new element to the array. + /// # Panic + /// This operation panics iff the length of all values (in bytes) exceeds `O` maximum value. pub fn push>(&mut self, value: Option) { self.try_push(value).unwrap() } @@ -176,3 +151,33 @@ impl> TryExtend> for MutableBinaryArray { iter.try_for_each(|x| self.try_push(x)) } } + +impl> TryPush> for MutableBinaryArray { + fn try_push(&mut self, value: Option) -> Result<()> { + match value { + Some(value) => { + let bytes = value.as_ref(); + + let size = O::from_usize(self.values.len() + bytes.len()) + .ok_or(ArrowError::KeyOverflowError)?; + + self.values.extend_from_slice(bytes); + + self.offsets.push(size); + + match &mut self.validity { + Some(validity) => validity.push(true), + None => {} + } + } + None => { + self.offsets.push(self.last_offset()); + match &mut self.validity { + Some(validity) => validity.push(false), + None => self.init_validity(), + } + } + } + Ok(()) + } +} diff --git a/src/array/boolean/mutable.rs b/src/array/boolean/mutable.rs index 603a30b83b7..32f3ba6d89b 100644 --- a/src/array/boolean/mutable.rs +++ b/src/array/boolean/mutable.rs @@ -1,9 +1,8 @@ use std::iter::FromIterator; use std::sync::Arc; -use crate::array::TryExtend; use crate::{ - array::{Array, MutableArray}, + array::{Array, MutableArray, TryExtend, TryPush}, bitmap::MutableBitmap, datatypes::DataType, error::Result, @@ -352,6 +351,14 @@ impl TryExtend> for MutableBooleanArray { } } +impl TryPush> for MutableBooleanArray { + /// This is infalible and is implemented for consistency with all other types + fn try_push(&mut self, item: Option) -> Result<()> { + self.push(item); + Ok(()) + } +} + impl PartialEq for MutableBooleanArray { fn eq(&self, other: &Self) -> bool { self.iter().eq(other.iter()) diff --git a/src/array/list/mutable.rs b/src/array/list/mutable.rs index 3dde5bede06..80a846b8e8d 100644 --- a/src/array/list/mutable.rs +++ b/src/array/list/mutable.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use crate::{ - array::{Array, MutableArray, Offset, TryExtend}, + array::{Array, MutableArray, Offset, TryExtend, TryPush}, bitmap::MutableBitmap, buffer::MutableBuffer, datatypes::{DataType, Field}, @@ -67,13 +67,25 @@ where { fn try_extend>>(&mut self, iter: II) -> Result<()> { for items in iter { - if let Some(items) = items { - let values = self.mut_values(); - values.try_extend(items)?; - self.try_push_valid()?; - } else { - self.push_null(); - } + self.try_push(items)?; + } + Ok(()) + } +} + +impl TryPush> for MutableListArray +where + O: Offset, + M: MutableArray + TryExtend>, + I: IntoIterator>, +{ + fn try_push(&mut self, item: Option) -> Result<()> { + if let Some(items) = item { + let values = self.mut_values(); + values.try_extend(items)?; + self.try_push_valid()?; + } else { + self.push_null(); } Ok(()) } diff --git a/src/array/mod.rs b/src/array/mod.rs index 2a9fcb3993a..09b1d3ac0de 100644 --- a/src/array/mod.rs +++ b/src/array/mod.rs @@ -414,6 +414,12 @@ pub trait TryExtend { fn try_extend>(&mut self, iter: I) -> Result<()>; } +/// A trait describing the ability of a struct to receive new items. +pub trait TryPush { + /// Tries to push a new element. + fn try_push(&mut self, item: A) -> Result<()>; +} + fn display_helper>>(iter: I) -> Vec { let iterator = iter.into_iter(); let len = iterator.size_hint().0; diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs index 19e352a34df..ee330188126 100644 --- a/src/array/primitive/mutable.rs +++ b/src/array/primitive/mutable.rs @@ -1,7 +1,7 @@ use std::{iter::FromIterator, sync::Arc}; use crate::{ - array::{Array, MutableArray, TryExtend}, + array::{Array, MutableArray, TryExtend, TryPush}, bitmap::MutableBitmap, buffer::MutableBuffer, datatypes::DataType, @@ -266,6 +266,14 @@ impl TryExtend> for MutablePrimitiveArray { } } +impl TryPush> for MutablePrimitiveArray { + /// This is infalible and is implemented for consistency with all other types + fn try_push(&mut self, item: Option) -> Result<()> { + self.push(item); + Ok(()) + } +} + impl MutableArray for MutablePrimitiveArray { fn len(&self) -> usize { self.values.len() diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs index 1e3b27020c6..61aa037aac5 100644 --- a/src/array/utf8/mutable.rs +++ b/src/array/utf8/mutable.rs @@ -3,7 +3,7 @@ use std::{iter::FromIterator, sync::Arc}; use crate::{ array::{ specification::{check_offsets, check_offsets_and_utf8}, - Array, MutableArray, Offset, TryExtend, + Array, MutableArray, Offset, TryExtend, TryPush, }, bitmap::MutableBitmap, buffer::MutableBuffer, @@ -126,35 +126,6 @@ impl MutableUtf8Array { *self.offsets.last().unwrap() } - /// Tries to push a new element to the array. - /// # Error - /// This operation errors iff the length of all values (in bytes) exceeds `O` maximum value. - pub fn try_push>(&mut self, value: Option) -> Result<()> { - match value { - Some(value) => { - let bytes = value.as_ref().as_bytes(); - self.values.extend_from_slice(bytes); - - let size = O::from_usize(self.values.len()).ok_or(ArrowError::KeyOverflowError)?; - - self.offsets.push(size); - - match &mut self.validity { - Some(validity) => validity.push(true), - None => {} - } - } - None => { - self.offsets.push(self.last_offset()); - match &mut self.validity { - Some(validity) => validity.push(false), - None => self.init_validity(), - } - } - } - Ok(()) - } - /// Pushes a new element to the array. /// # Panic /// This operation panics iff the length of all values (in bytes) exceeds `O` maximum value. @@ -340,6 +311,34 @@ impl> TryExtend> for MutableUtf8Array { } } +impl> TryPush> for MutableUtf8Array { + fn try_push(&mut self, value: Option) -> Result<()> { + match value { + Some(value) => { + let bytes = value.as_ref().as_bytes(); + self.values.extend_from_slice(bytes); + + let size = O::from_usize(self.values.len()).ok_or(ArrowError::KeyOverflowError)?; + + self.offsets.push(size); + + match &mut self.validity { + Some(validity) => validity.push(true), + None => {} + } + } + None => { + self.offsets.push(self.last_offset()); + match &mut self.validity { + Some(validity) => validity.push(false), + None => self.init_validity(), + } + } + } + Ok(()) + } +} + /// Creates [`MutableBitmap`] and two [`MutableBuffer`]s from an iterator of `Option`. /// The first buffer corresponds to a offset buffer, the second one /// corresponds to a values buffer. diff --git a/tests/it/array/list/mutable.rs b/tests/it/array/list/mutable.rs index d05ad6f34ca..2102735a935 100644 --- a/tests/it/array/list/mutable.rs +++ b/tests/it/array/list/mutable.rs @@ -29,3 +29,11 @@ fn basics() { ); assert_eq!(expected, array); } + +#[test] +fn push() { + let mut array = MutableListArray::>::new(); + array + .try_push(Some(vec![Some(1i32), Some(2), Some(3)])) + .unwrap(); +}