Into arrow with hint (#1730)

Author: joseph-isaacs

vortex-array/src/array/varbin/canonical.rs

@@ -28,6 +28,16 @@ impl IntoCanonical for VarBinArray {
         // Arrow representation.
         varbin_to_arrow(&self)
     }
+
+    fn into_arrow_with_data_type(self, data_type: &DataType) -> VortexResult<ArrayRef> {
+        let array_ref = self.into_arrow()?;
+
+        Ok(if array_ref.data_type() != data_type {
+            arrow_cast::cast(array_ref.as_ref(), data_type)?
+        } else {
+            array_ref
+        })
+    }
 }
 
 #[cfg(test)]

vortex-array/src/arrow/datum.rs

@@ -2,10 +2,10 @@ use arrow_array::{Array, ArrayRef, Datum as ArrowDatum};
 use vortex_error::VortexError;
 
 use crate::compute::slice;
-use crate::stats::{ArrayStatistics, Stat};
 use crate::{ArrayData, IntoCanonical};
 
 /// A wrapper around a generic Arrow array that can be used as a Datum in Arrow compute.
+#[derive(Debug)]
 pub struct Datum {
     array: ArrayRef,
     is_scalar: bool,
@@ -15,11 +15,7 @@ impl TryFrom<ArrayData> for Datum {
     type Error = VortexError;
 
     fn try_from(array: ArrayData) -> Result<Self, Self::Error> {
-        if array
-            .statistics()
-            .get_as::<bool>(Stat::IsConstant)
-            .unwrap_or_default()
-        {
+        if array.is_constant() {
             Ok(Self {
                 array: slice(array, 0, 1)?.into_arrow()?,
                 is_scalar: true,

vortex-array/src/arrow/record_batch.rs

@@ -1,5 +1,6 @@
-use arrow_array::cast::as_struct_array;
+use arrow_array::cast::AsArray;
 use arrow_array::RecordBatch;
+use arrow_schema::{DataType, Schema};
 use itertools::Itertools;
 use vortex_error::{vortex_err, VortexError, VortexResult};
 
@@ -41,16 +42,19 @@ impl TryFrom<ArrayData> for RecordBatch {
             vortex_err!("RecordBatch can only be constructed from a Vortex StructArray: {err}")
         })?;
 
-        RecordBatch::try_from(struct_arr)
+        struct_arr.into_record_batch()
     }
 }
 
-impl TryFrom<StructArray> for RecordBatch {
-    type Error = VortexError;
+impl StructArray {
+    pub fn into_record_batch(self) -> VortexResult<RecordBatch> {
+        let array_ref = self.into_array().into_arrow()?;
+        Ok(RecordBatch::from(array_ref.as_struct()))
+    }
 
-    fn try_from(value: StructArray) -> VortexResult<Self> {
-        let array_ref = value.into_canonical()?.into_arrow()?;
-        let struct_array = as_struct_array(array_ref.as_ref());
-        Ok(Self::from(struct_array))
+    pub fn into_record_batch_with_schema(self, schema: &Schema) -> VortexResult<RecordBatch> {
+        let data_type = DataType::Struct(schema.fields.clone());
+        let array_ref = self.into_array().into_arrow_with_data_type(&data_type)?;
+        Ok(RecordBatch::from(array_ref.as_struct()))
     }
 }