Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Added trait TryPush #314

Merged
merged 1 commit into from
Aug 22, 2021
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
63 changes: 34 additions & 29 deletions src/array/binary/mutable.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -61,34 +61,9 @@ impl<O: Offset> MutableBinaryArray<O> {
*self.offsets.last().unwrap()
}

pub fn try_push<T: AsRef<[u8]>>(&mut self, value: Option<T>) -> 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<T: AsRef<[u8]>>(&mut self, value: Option<T>) {
self.try_push(value).unwrap()
}
Expand Down Expand Up @@ -176,3 +151,33 @@ impl<O: Offset, T: AsRef<[u8]>> TryExtend<Option<T>> for MutableBinaryArray<O> {
iter.try_for_each(|x| self.try_push(x))
}
}

impl<O: Offset, T: AsRef<[u8]>> TryPush<Option<T>> for MutableBinaryArray<O> {
fn try_push(&mut self, value: Option<T>) -> 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(())
}
}
11 changes: 9 additions & 2 deletions src/array/boolean/mutable.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -352,6 +351,14 @@ impl TryExtend<Option<bool>> for MutableBooleanArray {
}
}

impl TryPush<Option<bool>> for MutableBooleanArray {
/// This is infalible and is implemented for consistency with all other types
fn try_push(&mut self, item: Option<bool>) -> Result<()> {
self.push(item);
Ok(())
}
}

impl PartialEq for MutableBooleanArray {
fn eq(&self, other: &Self) -> bool {
self.iter().eq(other.iter())
Expand Down
28 changes: 20 additions & 8 deletions src/array/list/mutable.rs
Original file line number Diff line number Diff line change
@@ -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},
Expand Down Expand Up @@ -67,13 +67,25 @@ where
{
fn try_extend<II: IntoIterator<Item = Option<I>>>(&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<O, M, I, T> TryPush<Option<I>> for MutableListArray<O, M>
where
O: Offset,
M: MutableArray + TryExtend<Option<T>>,
I: IntoIterator<Item = Option<T>>,
{
fn try_push(&mut self, item: Option<I>) -> Result<()> {
if let Some(items) = item {
let values = self.mut_values();
values.try_extend(items)?;
self.try_push_valid()?;
} else {
self.push_null();
}
Ok(())
}
Expand Down
6 changes: 6 additions & 0 deletions src/array/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -414,6 +414,12 @@ pub trait TryExtend<A> {
fn try_extend<I: IntoIterator<Item = A>>(&mut self, iter: I) -> Result<()>;
}

/// A trait describing the ability of a struct to receive new items.
pub trait TryPush<A> {
/// Tries to push a new element.
fn try_push(&mut self, item: A) -> Result<()>;
}

fn display_helper<T: std::fmt::Display, I: IntoIterator<Item = Option<T>>>(iter: I) -> Vec<String> {
let iterator = iter.into_iter();
let len = iterator.size_hint().0;
Expand Down
10 changes: 9 additions & 1 deletion src/array/primitive/mutable.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -266,6 +266,14 @@ impl<T: NativeType> TryExtend<Option<T>> for MutablePrimitiveArray<T> {
}
}

impl<T: NativeType> TryPush<Option<T>> for MutablePrimitiveArray<T> {
/// This is infalible and is implemented for consistency with all other types
fn try_push(&mut self, item: Option<T>) -> Result<()> {
self.push(item);
Ok(())
}
}

impl<T: NativeType> MutableArray for MutablePrimitiveArray<T> {
fn len(&self) -> usize {
self.values.len()
Expand Down
59 changes: 29 additions & 30 deletions src/array/utf8/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -126,35 +126,6 @@ impl<O: Offset> MutableUtf8Array<O> {
*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<T: AsRef<str>>(&mut self, value: Option<T>) -> 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.
Expand Down Expand Up @@ -340,6 +311,34 @@ impl<O: Offset, T: AsRef<str>> TryExtend<Option<T>> for MutableUtf8Array<O> {
}
}

impl<O: Offset, T: AsRef<str>> TryPush<Option<T>> for MutableUtf8Array<O> {
fn try_push(&mut self, value: Option<T>) -> 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.
Expand Down
8 changes: 8 additions & 0 deletions tests/it/array/list/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,3 +29,11 @@ fn basics() {
);
assert_eq!(expected, array);
}

#[test]
fn push() {
let mut array = MutableListArray::<i32, MutablePrimitiveArray<i32>>::new();
array
.try_push(Some(vec![Some(1i32), Some(2), Some(3)]))
.unwrap();
}