arrow-array/src/builder/primitive_run_builder.rs - arrow-rs - Git at Google

 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
 // regarding copyright ownership.  The ASF licenses this file
 // to you under the Apache License, Version 2.0 (the
 // "License"); you may not use this file except in compliance
 // with the License.  You may obtain a copy of the License at
 //
 //   http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing,
 // software distributed under the License is distributed on an
 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.

 use std::{any::Any, sync::Arc};

 use crate::{ArrayRef, ArrowPrimitiveType, RunArray, types::RunEndIndexType};

 use super::{ArrayBuilder, PrimitiveBuilder};

 use arrow_buffer::ArrowNativeType;

 /// Builder for [`RunArray`] of [`PrimitiveArray`](crate::array::PrimitiveArray)
 ///
 /// # Example:
 ///
 /// ```
 ///
 /// # use arrow_array::builder::PrimitiveRunBuilder;
 /// # use arrow_array::cast::AsArray;
 /// # use arrow_array::types::{UInt32Type, Int16Type};
 /// # use arrow_array::{Array, UInt32Array, Int16Array};
 ///
 /// let mut builder =
 /// PrimitiveRunBuilder::<Int16Type, UInt32Type>::new();
 /// builder.append_value(1234);
 /// builder.append_value(1234);
 /// builder.append_value(1234);
 /// builder.append_null();
 /// builder.append_value(5678);
 /// builder.append_value(5678);
 /// let array = builder.finish();
 ///
 /// assert_eq!(array.run_ends().values(), &[3, 4, 6]);
 ///
 /// let av = array.values();
 ///
 /// assert!(!av.is_null(0));
 /// assert!(av.is_null(1));
 /// assert!(!av.is_null(2));
 ///
 /// // Values are polymorphic and so require a downcast.
 /// let ava: &UInt32Array = av.as_primitive::<UInt32Type>();
 ///
 /// assert_eq!(ava, &UInt32Array::from(vec![Some(1234), None, Some(5678)]));
 /// ```
 #[derive(Debug)]
 pub struct PrimitiveRunBuilder<R, V>
 where
     R: RunEndIndexType,
     V: ArrowPrimitiveType,
 {
     run_ends_builder: PrimitiveBuilder<R>,
     values_builder: PrimitiveBuilder<V>,
     current_value: Option<V::Native>,
     current_run_end_index: usize,
     prev_run_end_index: usize,
 }

 impl<R, V> Default for PrimitiveRunBuilder<R, V>
 where
     R: RunEndIndexType,
     V: ArrowPrimitiveType,
 {
     fn default() -> Self {
         Self::new()
     }
 }

 impl<R, V> PrimitiveRunBuilder<R, V>
 where
     R: RunEndIndexType,
     V: ArrowPrimitiveType,
 {
     /// Creates a new `PrimitiveRunBuilder`
     pub fn new() -> Self {
         Self {
             run_ends_builder: PrimitiveBuilder::new(),
             values_builder: PrimitiveBuilder::new(),
             current_value: None,
             current_run_end_index: 0,
             prev_run_end_index: 0,
         }
     }

     /// Creates a new `PrimitiveRunBuilder` with the provided capacity
     ///
     /// `capacity`: the expected number of run-end encoded values.
     pub fn with_capacity(capacity: usize) -> Self {
         Self {
             run_ends_builder: PrimitiveBuilder::with_capacity(capacity),
             values_builder: PrimitiveBuilder::with_capacity(capacity),
             current_value: None,
             current_run_end_index: 0,
             prev_run_end_index: 0,
         }
     }
 }

 impl<R, V> ArrayBuilder for PrimitiveRunBuilder<R, V>
 where
     R: RunEndIndexType,
     V: ArrowPrimitiveType,
 {
     /// Returns the builder as a non-mutable `Any` reference.
     fn as_any(&self) -> &dyn Any {
         self
     }

     /// Returns the builder as a mutable `Any` reference.
     fn as_any_mut(&mut self) -> &mut dyn Any {
         self
     }

     /// Returns the boxed builder as a box of `Any`.
     fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
         self
     }

     /// Returns the length of logical array encoded by
     /// the eventual runs array.
     fn len(&self) -> usize {
         self.current_run_end_index
     }

     /// Builds the array and reset this builder.
     fn finish(&mut self) -> ArrayRef {
         Arc::new(self.finish())
     }

     /// Builds the array without resetting the builder.
     fn finish_cloned(&self) -> ArrayRef {
         Arc::new(self.finish_cloned())
     }
 }

 impl<R, V> PrimitiveRunBuilder<R, V>
 where
     R: RunEndIndexType,
     V: ArrowPrimitiveType,
 {
     /// Appends optional value to the logical array encoded by the RunArray.
     pub fn append_option(&mut self, value: Option<V::Native>) {
         if self.current_run_end_index == 0 {
             self.current_run_end_index = 1;
             self.current_value = value;
             return;
         }
         if self.current_value != value {
             self.append_run_end();
             self.current_value = value;
         }

         self.current_run_end_index += 1;
     }

     /// Appends value to the logical array encoded by the run-ends array.
     pub fn append_value(&mut self, value: V::Native) {
         self.append_option(Some(value))
     }

     /// Appends null to the logical array encoded by the run-ends array.
     pub fn append_null(&mut self) {
         self.append_option(None)
     }

     /// Creates the RunArray and resets the builder.
     /// Panics if RunArray cannot be built.
     pub fn finish(&mut self) -> RunArray<R> {
         // write the last run end to the array.
         self.append_run_end();

         // reset the run index to zero.
         self.current_value = None;
         self.current_run_end_index = 0;

         // build the run encoded array by adding run_ends and values array as its children.
         let run_ends_array = self.run_ends_builder.finish();
         let values_array = self.values_builder.finish();
         RunArray::<R>::try_new(&run_ends_array, &values_array).unwrap()
     }

     /// Creates the RunArray and without resetting the builder.
     /// Panics if RunArray cannot be built.
     pub fn finish_cloned(&self) -> RunArray<R> {
         let mut run_ends_array = self.run_ends_builder.finish_cloned();
         let mut values_array = self.values_builder.finish_cloned();

         // Add current run if one exists
         if self.prev_run_end_index != self.current_run_end_index {
             let mut run_end_builder = run_ends_array.into_builder().unwrap();
             let mut values_builder = values_array.into_builder().unwrap();
             self.append_run_end_with_builders(&mut run_end_builder, &mut values_builder);
             run_ends_array = run_end_builder.finish();
             values_array = values_builder.finish();
         }

         RunArray::try_new(&run_ends_array, &values_array).unwrap()
     }

     // Appends the current run to the array.
     fn append_run_end(&mut self) {
         // empty array or the function called without appending any value.
         if self.prev_run_end_index == self.current_run_end_index {
             return;
         }
         let run_end_index = self.run_end_index_as_native();
         self.run_ends_builder.append_value(run_end_index);
         self.values_builder.append_option(self.current_value);
         self.prev_run_end_index = self.current_run_end_index;
     }

     // Similar to `append_run_end` but on custom builders.
     // Used in `finish_cloned` which is not suppose to mutate `self`.
     fn append_run_end_with_builders(
         &self,
         run_ends_builder: &mut PrimitiveBuilder<R>,
         values_builder: &mut PrimitiveBuilder<V>,
     ) {
         let run_end_index = self.run_end_index_as_native();
         run_ends_builder.append_value(run_end_index);
         values_builder.append_option(self.current_value);
     }

     fn run_end_index_as_native(&self) -> R::Native {
         R::Native::from_usize(self.current_run_end_index)
         .unwrap_or_else(|| panic!(
                 "Cannot convert `current_run_end_index` {} from `usize` to native form of arrow datatype {}",
                 self.current_run_end_index,
                 R::DATA_TYPE
         ))
     }
 }

 impl<R, V> Extend<Option<V::Native>> for PrimitiveRunBuilder<R, V>
 where
     R: RunEndIndexType,
     V: ArrowPrimitiveType,
 {
     fn extend<T: IntoIterator<Item = Option<V::Native>>>(&mut self, iter: T) {
         for elem in iter {
             self.append_option(elem);
         }
     }
 }

 #[cfg(test)]
 mod tests {
     use crate::builder::PrimitiveRunBuilder;
     use crate::cast::AsArray;
     use crate::types::{Int16Type, UInt32Type};
     use crate::{Array, UInt32Array};

     #[test]
     fn test_primitive_ree_array_builder() {
         let mut builder = PrimitiveRunBuilder::<Int16Type, UInt32Type>::new();
         builder.append_value(1234);
         builder.append_value(1234);
         builder.append_value(1234);
         builder.append_null();
         builder.append_value(5678);
         builder.append_value(5678);

         let array = builder.finish();

         assert_eq!(array.null_count(), 0);
         assert_eq!(array.logical_null_count(), 1);
         assert_eq!(array.len(), 6);

         assert_eq!(array.run_ends().values(), &[3, 4, 6]);

         let av = array.values();

         assert!(!av.is_null(0));
         assert!(av.is_null(1));
         assert!(!av.is_null(2));

         // Values are polymorphic and so require a downcast.
         let ava: &UInt32Array = av.as_primitive::<UInt32Type>();

         assert_eq!(ava, &UInt32Array::from(vec![Some(1234), None, Some(5678)]));
     }

     #[test]
     fn test_extend() {
         let mut builder = PrimitiveRunBuilder::<Int16Type, Int16Type>::new();
         builder.extend([1, 2, 2, 5, 5, 4, 4].into_iter().map(Some));
         builder.extend([4, 4, 6, 2].into_iter().map(Some));
         let array = builder.finish();

         assert_eq!(array.len(), 11);
         assert_eq!(array.null_count(), 0);
         assert_eq!(array.logical_null_count(), 0);
         assert_eq!(array.run_ends().values(), &[1, 3, 5, 9, 10, 11]);
         assert_eq!(
             array.values().as_primitive::<Int16Type>().values(),
             &[1, 2, 5, 4, 6, 2]
         );
     }
 }
	// Licensed to the Apache Software Foundation (ASF) under one
	// or more contributor license agreements. See the NOTICE file
	// distributed with this work for additional information
	// regarding copyright ownership. The ASF licenses this file
	// to you under the Apache License, Version 2.0 (the
	// "License"); you may not use this file except in compliance
	// with the License. You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing,
	// software distributed under the License is distributed on an
	// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	// KIND, either express or implied. See the License for the
	// specific language governing permissions and limitations
	// under the License.

	use std::{any::Any, sync::Arc};

	use crate::{ArrayRef, ArrowPrimitiveType, RunArray, types::RunEndIndexType};

	use super::{ArrayBuilder, PrimitiveBuilder};

	use arrow_buffer::ArrowNativeType;

	/// Builder for [`RunArray`] of [`PrimitiveArray`](crate::array::PrimitiveArray)
	///
	/// # Example:
	///
	/// ```
	///
	/// # use arrow_array::builder::PrimitiveRunBuilder;
	/// # use arrow_array::cast::AsArray;
	/// # use arrow_array::types::{UInt32Type, Int16Type};
	/// # use arrow_array::{Array, UInt32Array, Int16Array};
	///
	/// let mut builder =
	/// PrimitiveRunBuilder::<Int16Type, UInt32Type>::new();
	/// builder.append_value(1234);
	/// builder.append_value(1234);
	/// builder.append_value(1234);
	/// builder.append_null();
	/// builder.append_value(5678);
	/// builder.append_value(5678);
	/// let array = builder.finish();
	///
	/// assert_eq!(array.run_ends().values(), &[3, 4, 6]);
	///
	/// let av = array.values();
	///
	/// assert!(!av.is_null(0));
	/// assert!(av.is_null(1));
	/// assert!(!av.is_null(2));
	///
	/// // Values are polymorphic and so require a downcast.
	/// let ava: &UInt32Array = av.as_primitive::<UInt32Type>();
	///
	/// assert_eq!(ava, &UInt32Array::from(vec![Some(1234), None, Some(5678)]));
	/// ```
	#[derive(Debug)]
	pub struct PrimitiveRunBuilder<R, V>
	where
	R: RunEndIndexType,
	V: ArrowPrimitiveType,
	{
	run_ends_builder: PrimitiveBuilder<R>,
	values_builder: PrimitiveBuilder<V>,
	current_value: Option<V::Native>,
	current_run_end_index: usize,
	prev_run_end_index: usize,
	}

	impl<R, V> Default for PrimitiveRunBuilder<R, V>
	where
	R: RunEndIndexType,
	V: ArrowPrimitiveType,
	{
	fn default() -> Self {
	Self::new()
	}
	}

	impl<R, V> PrimitiveRunBuilder<R, V>
	where
	R: RunEndIndexType,
	V: ArrowPrimitiveType,
	{
	/// Creates a new `PrimitiveRunBuilder`
	pub fn new() -> Self {
	Self {
	run_ends_builder: PrimitiveBuilder::new(),
	values_builder: PrimitiveBuilder::new(),
	current_value: None,
	current_run_end_index: 0,
	prev_run_end_index: 0,
	}
	}

	/// Creates a new `PrimitiveRunBuilder` with the provided capacity
	///
	/// `capacity`: the expected number of run-end encoded values.
	pub fn with_capacity(capacity: usize) -> Self {
	Self {
	run_ends_builder: PrimitiveBuilder::with_capacity(capacity),
	values_builder: PrimitiveBuilder::with_capacity(capacity),
	current_value: None,
	current_run_end_index: 0,
	prev_run_end_index: 0,
	}
	}
	}

	impl<R, V> ArrayBuilder for PrimitiveRunBuilder<R, V>
	where
	R: RunEndIndexType,
	V: ArrowPrimitiveType,
	{
	/// Returns the builder as a non-mutable `Any` reference.
	fn as_any(&self) -> &dyn Any {
	self
	}

	/// Returns the builder as a mutable `Any` reference.
	fn as_any_mut(&mut self) -> &mut dyn Any {
	self
	}

	/// Returns the boxed builder as a box of `Any`.
	fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
	self
	}

	/// Returns the length of logical array encoded by
	/// the eventual runs array.
	fn len(&self) -> usize {
	self.current_run_end_index
	}

	/// Builds the array and reset this builder.
	fn finish(&mut self) -> ArrayRef {
	Arc::new(self.finish())
	}

	/// Builds the array without resetting the builder.
	fn finish_cloned(&self) -> ArrayRef {
	Arc::new(self.finish_cloned())
	}
	}

	impl<R, V> PrimitiveRunBuilder<R, V>
	where
	R: RunEndIndexType,
	V: ArrowPrimitiveType,
	{
	/// Appends optional value to the logical array encoded by the RunArray.
	pub fn append_option(&mut self, value: Option<V::Native>) {
	if self.current_run_end_index == 0 {
	self.current_run_end_index = 1;
	self.current_value = value;
	return;
	}
	if self.current_value != value {
	self.append_run_end();
	self.current_value = value;
	}

	self.current_run_end_index += 1;
	}

	/// Appends value to the logical array encoded by the run-ends array.
	pub fn append_value(&mut self, value: V::Native) {
	self.append_option(Some(value))
	}

	/// Appends null to the logical array encoded by the run-ends array.
	pub fn append_null(&mut self) {
	self.append_option(None)
	}

	/// Creates the RunArray and resets the builder.
	/// Panics if RunArray cannot be built.
	pub fn finish(&mut self) -> RunArray<R> {
	// write the last run end to the array.
	self.append_run_end();

	// reset the run index to zero.
	self.current_value = None;
	self.current_run_end_index = 0;

	// build the run encoded array by adding run_ends and values array as its children.
	let run_ends_array = self.run_ends_builder.finish();
	let values_array = self.values_builder.finish();
	RunArray::<R>::try_new(&run_ends_array, &values_array).unwrap()
	}

	/// Creates the RunArray and without resetting the builder.
	/// Panics if RunArray cannot be built.
	pub fn finish_cloned(&self) -> RunArray<R> {
	let mut run_ends_array = self.run_ends_builder.finish_cloned();
	let mut values_array = self.values_builder.finish_cloned();

	// Add current run if one exists
	if self.prev_run_end_index != self.current_run_end_index {
	let mut run_end_builder = run_ends_array.into_builder().unwrap();
	let mut values_builder = values_array.into_builder().unwrap();
	self.append_run_end_with_builders(&mut run_end_builder, &mut values_builder);
	run_ends_array = run_end_builder.finish();
	values_array = values_builder.finish();
	}

	RunArray::try_new(&run_ends_array, &values_array).unwrap()
	}

	// Appends the current run to the array.
	fn append_run_end(&mut self) {
	// empty array or the function called without appending any value.
	if self.prev_run_end_index == self.current_run_end_index {
	return;
	}
	let run_end_index = self.run_end_index_as_native();
	self.run_ends_builder.append_value(run_end_index);
	self.values_builder.append_option(self.current_value);
	self.prev_run_end_index = self.current_run_end_index;
	}

	// Similar to `append_run_end` but on custom builders.
	// Used in `finish_cloned` which is not suppose to mutate `self`.
	fn append_run_end_with_builders(
	&self,
	run_ends_builder: &mut PrimitiveBuilder<R>,
	values_builder: &mut PrimitiveBuilder<V>,
	) {
	let run_end_index = self.run_end_index_as_native();
	run_ends_builder.append_value(run_end_index);
	values_builder.append_option(self.current_value);
	}

	fn run_end_index_as_native(&self) -> R::Native {
	R::Native::from_usize(self.current_run_end_index)
	.unwrap_or_else(\|\| panic!(
	"Cannot convert `current_run_end_index` {} from `usize` to native form of arrow datatype {}",
	self.current_run_end_index,
	R::DATA_TYPE
	))
	}
	}

	impl<R, V> Extend<Option<V::Native>> for PrimitiveRunBuilder<R, V>
	where
	R: RunEndIndexType,
	V: ArrowPrimitiveType,
	{
	fn extend<T: IntoIterator<Item = Option<V::Native>>>(&mut self, iter: T) {
	for elem in iter {
	self.append_option(elem);
	}
	}
	}

	#[cfg(test)]
	mod tests {
	use crate::builder::PrimitiveRunBuilder;
	use crate::cast::AsArray;
	use crate::types::{Int16Type, UInt32Type};
	use crate::{Array, UInt32Array};

	#[test]
	fn test_primitive_ree_array_builder() {
	let mut builder = PrimitiveRunBuilder::<Int16Type, UInt32Type>::new();
	builder.append_value(1234);
	builder.append_value(1234);
	builder.append_value(1234);
	builder.append_null();
	builder.append_value(5678);
	builder.append_value(5678);

	let array = builder.finish();

	assert_eq!(array.null_count(), 0);
	assert_eq!(array.logical_null_count(), 1);
	assert_eq!(array.len(), 6);

	assert_eq!(array.run_ends().values(), &[3, 4, 6]);

	let av = array.values();

	assert!(!av.is_null(0));
	assert!(av.is_null(1));
	assert!(!av.is_null(2));

	// Values are polymorphic and so require a downcast.
	let ava: &UInt32Array = av.as_primitive::<UInt32Type>();

	assert_eq!(ava, &UInt32Array::from(vec![Some(1234), None, Some(5678)]));
	}

	#[test]
	fn test_extend() {
	let mut builder = PrimitiveRunBuilder::<Int16Type, Int16Type>::new();
	builder.extend([1, 2, 2, 5, 5, 4, 4].into_iter().map(Some));
	builder.extend([4, 4, 6, 2].into_iter().map(Some));
	let array = builder.finish();

	assert_eq!(array.len(), 11);
	assert_eq!(array.null_count(), 0);
	assert_eq!(array.logical_null_count(), 0);
	assert_eq!(array.run_ends().values(), &[1, 3, 5, 9, 10, 11]);
	assert_eq!(
	array.values().as_primitive::<Int16Type>().values(),
	&[1, 2, 5, 4, 6, 2]
	);
	}
	}