Use Arc in Datatype to reduce memory overhead (#3)

See jorgecarleitao#1469

---------

Co-authored-by: Clement Rey <cr.rey.clement@gmail.com>

arrow-pyarrow-integration-testing/src/c_stream.rs

@@ -33,7 +33,7 @@ pub fn from_rust_iterator(py: Python) -> PyResult<PyObject> {
  // initialize an array
  let array = Int32Array::from(&[Some(2), None, Some(1), None]);
  let array = StructArray::new(
- DataType::Struct(vec![Field::new("a", array.data_type().clone(), true)]),
+ DataType::Struct(vec![Field::new("a", array.data_type().clone(), true)].into()),
  vec![array.boxed()],
  None,
  )

examples/csv_read_parallel.rs

@@ -1,5 +1,6 @@
 use crossbeam_channel::unbounded;
 
+use std::sync::Arc;
 use std::thread;
 use std::time::SystemTime;
 
@@ -18,7 +19,7 @@ fn parallel_read(path: &str) -> Result<Vec<Chunk<Box<dyn Array>>>> {
  let mut reader = read::ReaderBuilder::new().from_path(path)?;
  let (fields, _) =
  read::infer_schema(&mut reader, Some(batch_size * 10), has_header, &read::infer)?;
- let fields = Box::new(fields);
+ let fields = Arc::new(fields);
 
  let start = SystemTime::now();
  // spawn a thread to produce `Vec<ByteRecords>` (IO bounded)

examples/extension.rs

@@ -1,4 +1,7 @@
-use std::io::{Cursor, Seek, Write};
+use std::{
+ io::{Cursor, Seek, Write},
+ sync::Arc,
+};
 
 use re_arrow2::array::*;
 use re_arrow2::chunk::Chunk;
@@ -10,7 +13,7 @@ use re_arrow2::io::ipc::write;
 fn main() -> Result<()> {
  // declare an extension.
  let extension_type =
- DataType::Extension("date16".to_string(), Box::new(DataType::UInt16), None);
+ DataType::Extension("date16".to_string(), Arc::new(DataType::UInt16), None);
 
  // initialize an array with it.
  let array = UInt16Array::from_slice([1, 2]).to(extension_type.clone());

src/array/dictionary/mod.rs

@@ -1,5 +1,4 @@
-use std::hash::Hash;
-use std::hint::unreachable_unchecked;
+use std::{hash::Hash, hint::unreachable_unchecked, sync::Arc};
 
 use crate::{
  bitmap::{
@@ -292,7 +291,7 @@ impl<K: DictionaryKey> DictionaryArray<K> {
  }
 
  pub(crate) fn default_data_type(values_datatype: DataType) -> DataType {
- DataType::Dictionary(K::KEY_TYPE, Box::new(values_datatype), false)
+ DataType::Dictionary(K::KEY_TYPE, Arc::new(values_datatype), false)
  }
 
  /// Slices this [`DictionaryArray`].

src/array/dictionary/mutable.rs

@@ -86,7 +86,7 @@ impl<K: DictionaryKey, M: MutableArray> MutableDictionaryArray<K, M> {
  fn from_value_map(value_map: ValueMap<K, M>) -> Self {
  let keys = MutablePrimitiveArray::<K>::new();
  let data_type =
- DataType::Dictionary(K::KEY_TYPE, Box::new(value_map.data_type().clone()), false);
+ DataType::Dictionary(K::KEY_TYPE, Arc::new(value_map.data_type().clone()), false);
  Self {
  data_type,
  map: value_map,

src/array/fixed_size_list/mod.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use crate::{
  bitmap::Bitmap,
  datatypes::{DataType, Field},
@@ -203,7 +205,7 @@ impl FixedSizeListArray {
 
  /// Returns a [`DataType`] consistent with [`FixedSizeListArray`].
  pub fn default_datatype(data_type: DataType, size: usize) -> DataType {
- let field = Box::new(Field::new("item", data_type, true));
+ let field = Arc::new(Field::new("item", data_type, true));
  DataType::FixedSizeList(field, size)
  }
 }

src/array/fixed_size_list/mutable.rs

@@ -42,7 +42,7 @@ impl<M: MutableArray> MutableFixedSizeListArray<M> {
  /// Creates a new [`MutableFixedSizeListArray`] from a [`MutableArray`] and size.
  pub fn new_with_field(values: M, name: &str, nullable: bool, size: usize) -> Self {
  let data_type = DataType::FixedSizeList(
- Box::new(Field::new(name, values.data_type().clone(), nullable)),
+ Arc::new(Field::new(name, values.data_type().clone(), nullable)),
  size,
  );
  Self::new_from(values, data_type, size)

src/array/list/mod.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use crate::{
  bitmap::Bitmap,
  datatypes::{DataType, Field},
@@ -188,7 +190,7 @@ impl<O: Offset> ListArray<O> {
 impl<O: Offset> ListArray<O> {
  /// Returns a default [`DataType`]: inner field is named "item" and is nullable
  pub fn default_datatype(data_type: DataType) -> DataType {
- let field = Box::new(Field::new("item", data_type, true));
+ let field = Arc::new(Field::new("item", data_type, true));
  if O::IS_LARGE {
  DataType::LargeList(field)
  } else {

src/array/list/mutable.rs

@@ -127,7 +127,7 @@ impl<O: Offset, M: MutableArray> MutableListArray<O, M> {
 
  /// Creates a new [`MutableListArray`] from a [`MutableArray`].
  pub fn new_with_field(values: M, name: &str, nullable: bool) -> Self {
- let field = Box::new(Field::new(name, values.data_type().clone(), nullable));
+ let field = Arc::new(Field::new(name, values.data_type().clone(), nullable));
  let data_type = if O::IS_LARGE {
  DataType::LargeList(field)
  } else {

src/array/struct_/mod.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use crate::{
  bitmap::Bitmap,
  datatypes::{DataType, Field},
@@ -28,7 +30,7 @@ pub use mutable::*;
 /// Field::new("c", DataType::Int32, false),
 /// ];
 ///
-/// let array = StructArray::new(DataType::Struct(fields), vec![boolean, int], None);
+/// let array = StructArray::new(DataType::Struct(std::sync::Arc::new(fields)), vec![boolean, int], None);
 /// ```
 #[derive(Clone)]
 pub struct StructArray {
@@ -153,7 +155,7 @@ impl StructArray {
 impl StructArray {
  /// Deconstructs the [`StructArray`] into its individual components.
  #[must_use]
- pub fn into_data(self) -> (Vec<Field>, Vec<Box<dyn Array>>, Option<Bitmap>) {
+ pub fn into_data(self) -> (Arc<Vec<Field>>, Vec<Box<dyn Array>>, Option<Bitmap>) {
  let Self {
  data_type,
  values,

src/array/union/mod.rs

@@ -73,7 +73,7 @@ impl UnionArray {
  .try_for_each(|(index, (data_type, child))| {
  if data_type != child {
  Err(Error::oos(format!(
- "The children DataTypes of a UnionArray must equal the children data types. 
+ "The children DataTypes of a UnionArray must equal the children data types.
  However, the field {index} has data type {data_type:?} but the value has data type {child:?}"
  )))
  } else {
@@ -352,7 +352,7 @@ impl UnionArray {
  fn try_get_all(data_type: &DataType) -> Result<UnionComponents, Error> {
  match data_type.to_logical_type() {
  DataType::Union(fields, ids, mode) => {
- Ok((fields, ids.as_ref().map(|x| x.as_ref()), *mode))
+ Ok((fields, ids.as_ref().map(|x| x.as_slice()), *mode))
  }
  _ => Err(Error::oos(
  "The UnionArray requires a logical type of DataType::Union",

src/compute/arithmetics/time.rs

@@ -81,7 +81,7 @@ fn create_scale(lhs: &DataType, rhs: &DataType) -> Result<f64> {
 /// ])
 /// .to(DataType::Timestamp(
 /// TimeUnit::Second,
-/// Some("America/New_York".to_string()),
+/// Some(std::sync::Arc::new("America/New_york".to_string())),
 /// ));
 ///
 /// let duration = PrimitiveArray::from([Some(10i64), Some(20i64), None, Some(30i64)])
@@ -96,7 +96,7 @@ fn create_scale(lhs: &DataType, rhs: &DataType) -> Result<f64> {
 /// ])
 /// .to(DataType::Timestamp(
 /// TimeUnit::Second,
-/// Some("America/New_York".to_string()),
+/// Some(std::sync::Arc::new("America/New_york".to_string())),
 /// ));
 ///
 /// assert_eq!(result, expected);
@@ -161,7 +161,7 @@ where
 /// ])
 /// .to(DataType::Timestamp(
 /// TimeUnit::Second,
-/// Some("America/New_York".to_string()),
+/// Some(std::sync::Arc::new("America/New_york".to_string())),
 /// ));
 ///
 /// let duration = PrimitiveArray::from([Some(10i64), Some(20i64), None, Some(30i64)])
@@ -176,7 +176,7 @@ where
 /// ])
 /// .to(DataType::Timestamp(
 /// TimeUnit::Second,
-/// Some("America/New_York".to_string()),
+/// Some(std::sync::Arc::new("America/New_york".to_string())),
 /// ));
 ///
 /// assert_eq!(result, expected);

src/compute/cast/dictionary_to.rs

@@ -89,7 +89,7 @@ where
  } else {
  let data_type = DataType::Dictionary(
  K2::KEY_TYPE,
- Box::new(values.data_type().clone()),
+ std::sync::Arc::new(values.data_type().clone()),
  is_ordered,
  );
  // Safety: this is safe because given a type `T` that fits in a `usize`, casting it to type `P` either overflows or also fits in a `usize`
@@ -116,7 +116,7 @@ where
  } else {
  let data_type = DataType::Dictionary(
  K2::KEY_TYPE,
- Box::new(values.data_type().clone()),
+ std::sync::Arc::new(values.data_type().clone()),
  is_ordered,
  );
  // some of the values may not fit in `usize` and thus this needs to be checked

src/compute/cast/mod.rs

@@ -409,7 +409,7 @@ fn cast_list_to_fixed_size_list<O: Offset>(
  );
  let new_values = cast(sliced_values.as_ref(), inner.data_type(), options)?;
  Ok(FixedSizeListArray::new(
- DataType::FixedSizeList(Box::new(inner.clone()), size),
+ DataType::FixedSizeList(std::sync::Arc::new(inner.clone()), size),
  new_values,
  list.validity().cloned(),
  ))

src/compute/cast/primitive_to.rs

@@ -1,4 +1,5 @@
 use std::hash::Hash;
+use std::sync::Arc;
 
 use num_traits::{AsPrimitive, Float, ToPrimitive};
 
@@ -406,7 +407,7 @@ pub fn timestamp_to_timestamp(
  from: &PrimitiveArray<i64>,
  from_unit: TimeUnit,
  to_unit: TimeUnit,
- tz: &Option<String>,
+ tz: &Option<Arc<String>>,
 ) -> PrimitiveArray<i64> {
  let from_size = time_unit_multiple(from_unit);
  let to_size = time_unit_multiple(to_unit);

src/compute/cast/utf8_to.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use chrono::Datelike;
 
 use crate::{
@@ -127,7 +129,7 @@ pub fn utf8_to_naive_timestamp_ns<O: Offset>(from: &Utf8Array<O>) -> PrimitiveAr
 
 pub(super) fn utf8_to_timestamp_ns_dyn<O: Offset>(
  from: &dyn Array,
- timezone: String,
+ timezone: Arc<String>,
 ) -> Result<Box<dyn Array>> {
  let from = from.as_any().downcast_ref().unwrap();
  utf8_to_timestamp_ns::<O>(from, timezone)
@@ -138,7 +140,7 @@ pub(super) fn utf8_to_timestamp_ns_dyn<O: Offset>(
 /// [`crate::temporal_conversions::utf8_to_timestamp_ns`] applied for RFC3339 formatting
 pub fn utf8_to_timestamp_ns<O: Offset>(
  from: &Utf8Array<O>,
- timezone: String,
+ timezone: Arc<String>,
 ) -> Result<PrimitiveArray<i64>> {
  utf8_to_timestamp_ns_(from, RFC3339, timezone)
 }