add Array::as_any_arc

Closes apache#2901.
crepererum · Oct 20, 2022 · a483a0c · a483a0c
1 parent a881083
commit a483a0c
Show file tree

Hide file tree

Showing 13 changed files with 96 additions and 3 deletions.
diff --git a/arrow-array/src/array/binary_array.rs b/arrow-array/src/array/binary_array.rs
@@ -25,6 +25,7 @@ use arrow_buffer::{bit_util, Buffer, MutableBuffer};
 use arrow_data::ArrayData;
 use arrow_schema::DataType;
 use std::any::Any;
+use std::sync::Arc;
 
 /// See [`BinaryArray`] and [`LargeBinaryArray`] for storing
 /// binary data.
@@ -254,6 +255,10 @@ impl<OffsetSize: OffsetSizeTrait> Array for GenericBinaryArray<OffsetSize> {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }

diff --git a/arrow-array/src/array/boolean_array.rs b/arrow-array/src/array/boolean_array.rs
@@ -23,6 +23,7 @@ use arrow_buffer::{bit_util, Buffer, MutableBuffer};
 use arrow_data::ArrayData;
 use arrow_schema::DataType;
 use std::any::Any;
+use std::sync::Arc;
 
 /// Array of bools
 ///
@@ -152,6 +153,10 @@ impl Array for BooleanArray {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }

diff --git a/arrow-array/src/array/dictionary_array.rs b/arrow-array/src/array/dictionary_array.rs
@@ -26,6 +26,7 @@ use arrow_buffer::ArrowNativeType;
 use arrow_data::ArrayData;
 use arrow_schema::{ArrowError, DataType};
 use std::any::Any;
+use std::sync::Arc;
 
 ///
 /// A dictionary array where each element is a single value indexed by an integer key.
@@ -520,6 +521,10 @@ impl<T: ArrowPrimitiveType> Array for DictionaryArray<T> {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }
@@ -597,6 +602,10 @@ impl<'a, K: ArrowPrimitiveType, V: Sync> Array for TypedDictionaryArray<'a, K, V
         self.dictionary
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        None
+    }
+
     fn data(&self) -> &ArrayData {
         &self.dictionary.data
     }

diff --git a/arrow-array/src/array/fixed_size_binary_array.rs b/arrow-array/src/array/fixed_size_binary_array.rs
@@ -22,6 +22,7 @@ use arrow_buffer::{bit_util, Buffer, MutableBuffer};
 use arrow_data::ArrayData;
 use arrow_schema::{ArrowError, DataType};
 use std::any::Any;
+use std::sync::Arc;
 
 /// An array where each element is a fixed-size sequence of bytes.
 ///
@@ -358,6 +359,10 @@ impl Array for FixedSizeBinaryArray {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }

diff --git a/arrow-array/src/array/fixed_size_list_array.rs b/arrow-array/src/array/fixed_size_list_array.rs
@@ -22,6 +22,7 @@ use crate::{
 use arrow_data::ArrayData;
 use arrow_schema::DataType;
 use std::any::Any;
+use std::sync::Arc;
 
 /// A list array where each element is a fixed-size sequence of values with the same
 /// type whose maximum length is represented by a i32.
@@ -202,6 +203,10 @@ impl Array for FixedSizeListArray {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }

diff --git a/arrow-array/src/array/list_array.rs b/arrow-array/src/array/list_array.rs
@@ -26,6 +26,7 @@ use arrow_data::ArrayData;
 use arrow_schema::{ArrowError, DataType, Field};
 use num::Integer;
 use std::any::Any;
+use std::sync::Arc;
 
 /// trait declaring an offset size, relevant for i32 vs i64 array types.
 pub trait OffsetSizeTrait: ArrowNativeType + std::ops::AddAssign + Integer {
@@ -252,6 +253,10 @@ impl<OffsetSize: OffsetSizeTrait> Array for GenericListArray<OffsetSize> {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }

diff --git a/arrow-array/src/array/map_array.rs b/arrow-array/src/array/map_array.rs
@@ -208,6 +208,10 @@ impl Array for MapArray {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }

diff --git a/arrow-array/src/array/mod.rs b/arrow-array/src/array/mod.rs
@@ -88,6 +88,31 @@ pub trait Array: std::fmt::Debug + Send + Sync {
     /// ```
     fn as_any(&self) -> &dyn Any;
 
+    /// Returns the array as [`Any`](std::any::Any) wrapped into an [`Arc`] so that it can be
+    /// downcasted to a specific implementation without the need for an unsized reference. This is helpful if you want
+    /// to return a concrete array type.
+    /// 
+    /// Note that this conversion may not be possible for all types due to the `'static` constraint.
+    ///
+    /// # Example:
+    ///
+    /// ```
+    /// # use std::sync::Arc;
+    /// # use arrow_array::{Int32Array, RecordBatch};
+    /// # use arrow_schema::{Schema, Field, DataType, ArrowError};
+    ///
+    /// let id = Int32Array::from(vec![1, 2, 3, 4, 5]);
+    /// let batch = RecordBatch::try_new(
+    ///     Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])),
+    ///     vec![Arc::new(id)]
+    /// ).unwrap();
+    ///
+    /// let int32array = Arc::downcast::<Int32Array>(
+    ///     Arc::clone(batch.column(0)).as_any_arc().expect("Failed to create static Arc")
+    /// ).expect("Failed to downcast");
+    /// ```
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>>;
+
     /// Returns a reference to the underlying data of this array.
     fn data(&self) -> &ArrayData;
 
@@ -258,6 +283,10 @@ impl Array for ArrayRef {
         self.as_ref().as_any()
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         self.as_ref().data()
     }
@@ -316,6 +345,10 @@ impl<'a, T: Array> Array for &'a T {
         T::as_any(self)
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        None
+    }
+
     fn data(&self) -> &ArrayData {
         T::data(self)
     }

diff --git a/arrow-array/src/array/null_array.rs b/arrow-array/src/array/null_array.rs
@@ -20,7 +20,7 @@
 use crate::Array;
 use arrow_data::ArrayData;
 use arrow_schema::DataType;
-use std::any::Any;
+use std::{any::Any, sync::Arc};
 
 /// An Array where all elements are nulls
 ///
@@ -58,6 +58,10 @@ impl Array for NullArray {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }

diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs
@@ -29,6 +29,7 @@ use arrow_schema::{ArrowError, DataType};
 use chrono::{Duration, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime};
 use half::f16;
 use std::any::Any;
+use std::sync::Arc;
 
 ///
 /// # Example: Using `collect`
@@ -466,6 +467,10 @@ impl<T: ArrowPrimitiveType> Array for PrimitiveArray<T> {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }

diff --git a/arrow-array/src/array/string_array.rs b/arrow-array/src/array/string_array.rs
@@ -25,6 +25,7 @@ use arrow_buffer::{bit_util, Buffer, MutableBuffer};
 use arrow_data::ArrayData;
 use arrow_schema::DataType;
 use std::any::Any;
+use std::sync::Arc;
 
 /// Generic struct for \[Large\]StringArray
 ///
@@ -318,6 +319,10 @@ impl<OffsetSize: OffsetSizeTrait> Array for GenericStringArray<OffsetSize> {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }

diff --git a/arrow-array/src/array/struct_array.rs b/arrow-array/src/array/struct_array.rs
@@ -20,7 +20,7 @@ use arrow_buffer::buffer::buffer_bin_or;
 use arrow_buffer::Buffer;
 use arrow_data::ArrayData;
 use arrow_schema::{ArrowError, DataType, Field};
-use std::any::Any;
+use std::{any::Any, sync::Arc};
 
 /// A nested array type where each child (called *field*) is represented by a separate
 /// array.
@@ -192,6 +192,10 @@ impl Array for StructArray {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }

diff --git a/arrow-array/src/array/union_array.rs b/arrow-array/src/array/union_array.rs
@@ -21,7 +21,7 @@ use arrow_data::ArrayData;
 use arrow_schema::{ArrowError, DataType, Field, UnionMode};
 /// Contains the `UnionArray` type.
 ///
-use std::any::Any;
+use std::{any::Any, sync::Arc};
 
 /// An Array that can represent slots of varying types.
 ///
@@ -310,6 +310,10 @@ impl Array for UnionArray {
         self
     }
 
+    fn as_any_arc(self: Arc<Self>) -> Option<Arc<dyn Any + Send + Sync + 'static>> {
+        Some(self)
+    }
+
     fn data(&self) -> &ArrayData {
         &self.data
     }