pluggable registry for input/export arrow kernels

vortex-array/Cargo.toml

@@ -29,7 +29,7 @@ arrow-buffer = { workspace = true }
 arrow-cast = { workspace = true }
 arrow-data = { workspace = true }
 arrow-ord = { workspace = true }
-arrow-schema = { workspace = true }
+arrow-schema = { workspace = true, features = ["canonical_extension_types"] }
 arrow-select = { workspace = true }
 arrow-string = { workspace = true }
 async-lock = { workspace = true }

vortex-array/public-api.lock

@@ -7076,6 +7076,18 @@ pub mod vortex_array::arrow::primitive
 
 pub fn vortex_array::arrow::primitive::canonical_primitive_to_arrow<T: arrow_array::types::ArrowPrimitiveType>(vortex_array::arrays::PrimitiveArray, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<arrow_array::array::ArrayRef> where <T as arrow_array::types::ArrowPrimitiveType>::Native: vortex_array::dtype::NativePType
 
+pub enum vortex_array::arrow::ArrowExport
+
+pub vortex_array::arrow::ArrowExport::Exported(arrow_array::array::ArrayRef)
+
+pub vortex_array::arrow::ArrowExport::Unsupported(vortex_array::ArrayRef)
+
+pub enum vortex_array::arrow::ArrowImport
+
+pub vortex_array::arrow::ArrowImport::Imported(vortex_array::ArrayRef)
+
+pub vortex_array::arrow::ArrowImport::Unsupported(arrow_array::array::ArrayRef)
+
 pub struct vortex_array::arrow::ArrowArrayStreamAdapter
 
 impl vortex_array::arrow::ArrowArrayStreamAdapter
@@ -7092,6 +7104,44 @@ impl vortex_array::iter::ArrayIterator for vortex_array::arrow::ArrowArrayStream
 
 pub fn vortex_array::arrow::ArrowArrayStreamAdapter::dtype(&self) -> &vortex_array::dtype::DType
 
+pub struct vortex_array::arrow::ArrowSession
+
+impl vortex_array::arrow::ArrowSession
+
+pub fn vortex_array::arrow::ArrowSession::execute_arrow(&self, vortex_array::ArrayRef, core::option::Option<&arrow_schema::field::Field>, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<arrow_array::array::ArrayRef>
+
+pub fn vortex_array::arrow::ArrowSession::from_arrow_array(&self, arrow_array::array::ArrayRef, &arrow_schema::field::Field) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
+pub fn vortex_array::arrow::ArrowSession::from_arrow_field(&self, &arrow_schema::field::Field) -> vortex_error::VortexResult<vortex_array::dtype::DType>
+
+pub fn vortex_array::arrow::ArrowSession::from_arrow_record_batch(&self, arrow_array::record_batch::RecordBatch, &arrow_schema::schema::Schema) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+
+pub fn vortex_array::arrow::ArrowSession::from_arrow_schema(&self, &arrow_schema::schema::Schema) -> vortex_error::VortexResult<vortex_array::dtype::DType>
+
+pub fn vortex_array::arrow::ArrowSession::register_exporter(&self, vortex_array::arrow::ArrowExportVTableRef)
+
+pub fn vortex_array::arrow::ArrowSession::register_importer(&self, vortex_array::arrow::ArrowImportVTableRef)
+
+pub fn vortex_array::arrow::ArrowSession::to_arrow_data_type(&self, &vortex_array::dtype::DType) -> vortex_error::VortexResult<arrow_schema::datatype::DataType>
+
+pub fn vortex_array::arrow::ArrowSession::to_arrow_field(&self, &str, &vortex_array::dtype::DType) -> vortex_error::VortexResult<arrow_schema::field::Field>
+
+pub fn vortex_array::arrow::ArrowSession::to_arrow_schema(&self, &vortex_array::dtype::DType) -> vortex_error::VortexResult<arrow_schema::schema::Schema>
+
+impl core::default::Default for vortex_array::arrow::ArrowSession
+
+pub fn vortex_array::arrow::ArrowSession::default() -> Self
+
+impl core::fmt::Debug for vortex_array::arrow::ArrowSession
+
+pub fn vortex_array::arrow::ArrowSession::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_session::SessionVar for vortex_array::arrow::ArrowSession
+
+pub fn vortex_array::arrow::ArrowSession::as_any(&self) -> &dyn core::any::Any
+
+pub fn vortex_array::arrow::ArrowSession::as_any_mut(&mut self) -> &mut dyn core::any::Any
+
 pub struct vortex_array::arrow::Datum
 
 impl vortex_array::arrow::Datum
@@ -7128,6 +7178,50 @@ pub fn vortex_array::ArrayRef::execute_record_batch(self, &arrow_schema::schema:
 
 pub fn vortex_array::ArrayRef::execute_record_batches(self, &arrow_schema::schema::Schema, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<alloc::vec::Vec<arrow_array::record_batch::RecordBatch>>
 
+pub trait vortex_array::arrow::ArrowExportVTable: 'static + core::marker::Send + core::marker::Sync + core::fmt::Debug
+
+pub fn vortex_array::arrow::ArrowExportVTable::arrow_ext_id(&self) -> vortex_session::registry::Id
+
+pub fn vortex_array::arrow::ArrowExportVTable::execute_arrow(&self, vortex_array::ArrayRef, &arrow_schema::field::Field, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::arrow::ArrowExport>
+
+pub fn vortex_array::arrow::ArrowExportVTable::to_arrow_field(&self, &str, &vortex_array::dtype::extension::ExtDTypeRef) -> vortex_error::VortexResult<core::option::Option<arrow_schema::field::Field>>
+
+pub fn vortex_array::arrow::ArrowExportVTable::vortex_ext_id(&self) -> vortex_array::dtype::extension::ExtId
+
+impl vortex_array::arrow::ArrowExportVTable for vortex_array::extension::uuid::Uuid
+
+pub fn vortex_array::extension::uuid::Uuid::arrow_ext_id(&self) -> vortex_session::registry::Id
+
+pub fn vortex_array::extension::uuid::Uuid::execute_arrow(&self, vortex_array::ArrayRef, &arrow_schema::field::Field, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::arrow::ArrowExport>
+
+pub fn vortex_array::extension::uuid::Uuid::to_arrow_field(&self, &str, &vortex_array::dtype::extension::ExtDTypeRef) -> vortex_error::VortexResult<core::option::Option<arrow_schema::field::Field>>
+
+pub fn vortex_array::extension::uuid::Uuid::vortex_ext_id(&self) -> vortex_array::dtype::extension::ExtId
+
+pub trait vortex_array::arrow::ArrowImportVTable: 'static + core::marker::Send + core::marker::Sync + core::fmt::Debug
+
+pub fn vortex_array::arrow::ArrowImportVTable::arrow_ext_id(&self) -> vortex_session::registry::Id
+
+pub fn vortex_array::arrow::ArrowImportVTable::from_arrow_array(&self, arrow_array::array::ArrayRef, &vortex_array::dtype::extension::ExtDTypeRef) -> vortex_error::VortexResult<vortex_array::arrow::ArrowImport>
+
+pub fn vortex_array::arrow::ArrowImportVTable::from_arrow_field(&self, &arrow_schema::field::Field) -> vortex_error::VortexResult<core::option::Option<vortex_array::dtype::DType>>
+
+impl vortex_array::arrow::ArrowImportVTable for vortex_array::extension::uuid::Uuid
+
+pub fn vortex_array::extension::uuid::Uuid::arrow_ext_id(&self) -> vortex_session::registry::Id
+
+pub fn vortex_array::extension::uuid::Uuid::from_arrow_array(&self, arrow_array::array::ArrayRef, &vortex_array::dtype::extension::ExtDTypeRef) -> vortex_error::VortexResult<vortex_array::arrow::ArrowImport>
+
+pub fn vortex_array::extension::uuid::Uuid::from_arrow_field(&self, &arrow_schema::field::Field) -> vortex_error::VortexResult<core::option::Option<vortex_array::dtype::DType>>
+
+pub trait vortex_array::arrow::ArrowSessionExt: vortex_session::SessionExt
+
+pub fn vortex_array::arrow::ArrowSessionExt::arrow(&self) -> vortex_session::Ref<'_, vortex_array::arrow::ArrowSession>
+
+impl<S: vortex_session::SessionExt> vortex_array::arrow::ArrowSessionExt for S
+
+pub fn S::arrow(&self) -> vortex_session::Ref<'_, vortex_array::arrow::ArrowSession>
+
 pub trait vortex_array::arrow::FromArrowArray<A>
 
 pub fn vortex_array::arrow::FromArrowArray::from_arrow(A, bool) -> vortex_error::VortexResult<Self> where Self: core::marker::Sized
@@ -7298,6 +7392,10 @@ pub fn vortex_array::arrow::to_arrow_null_buffer(vortex_array::validity::Validit
 
 pub fn vortex_array::arrow::to_null_buffer(vortex_mask::Mask) -> core::option::Option<arrow_buffer::buffer::null::NullBuffer>
 
+pub type vortex_array::arrow::ArrowExportVTableRef = alloc::sync::Arc<dyn vortex_array::arrow::ArrowExportVTable>
+
+pub type vortex_array::arrow::ArrowImportVTableRef = alloc::sync::Arc<dyn vortex_array::arrow::ArrowImportVTable>
+
 pub mod vortex_array::buffer
 
 pub struct vortex_array::buffer::BufferHandle(_)
@@ -13118,6 +13216,24 @@ pub fn vortex_array::extension::uuid::Uuid::hash<__H: core::hash::Hasher>(&self,
 
 impl core::marker::StructuralPartialEq for vortex_array::extension::uuid::Uuid
 
+impl vortex_array::arrow::ArrowExportVTable for vortex_array::extension::uuid::Uuid
+
+pub fn vortex_array::extension::uuid::Uuid::arrow_ext_id(&self) -> vortex_session::registry::Id
+
+pub fn vortex_array::extension::uuid::Uuid::execute_arrow(&self, vortex_array::ArrayRef, &arrow_schema::field::Field, &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::arrow::ArrowExport>
+
+pub fn vortex_array::extension::uuid::Uuid::to_arrow_field(&self, &str, &vortex_array::dtype::extension::ExtDTypeRef) -> vortex_error::VortexResult<core::option::Option<arrow_schema::field::Field>>
+
+pub fn vortex_array::extension::uuid::Uuid::vortex_ext_id(&self) -> vortex_array::dtype::extension::ExtId
+
+impl vortex_array::arrow::ArrowImportVTable for vortex_array::extension::uuid::Uuid
+
+pub fn vortex_array::extension::uuid::Uuid::arrow_ext_id(&self) -> vortex_session::registry::Id
+
+pub fn vortex_array::extension::uuid::Uuid::from_arrow_array(&self, arrow_array::array::ArrayRef, &vortex_array::dtype::extension::ExtDTypeRef) -> vortex_error::VortexResult<vortex_array::arrow::ArrowImport>
+
+pub fn vortex_array::extension::uuid::Uuid::from_arrow_field(&self, &arrow_schema::field::Field) -> vortex_error::VortexResult<core::option::Option<vortex_array::dtype::DType>>
+
 impl vortex_array::dtype::extension::ExtVTable for vortex_array::extension::uuid::Uuid
 
 pub type vortex_array::extension::uuid::Uuid::Metadata = vortex_array::extension::uuid::UuidMetadata

vortex-array/src/arrow/convert.rs

@@ -331,7 +331,7 @@ impl FromArrowArray<&ArrowBooleanArray> for ArrayRef {
 }
 
 /// Strip out the nulls from this array and return a new array without nulls.
-fn remove_nulls(data: arrow_data::ArrayData) -> arrow_data::ArrayData {
+pub(crate) fn remove_nulls(data: arrow_data::ArrayData) -> arrow_data::ArrayData {
     if data.null_count() == 0 {
         // No nulls to remove, return the array as is
         return data;
@@ -476,7 +476,7 @@ impl<K: ArrowDictionaryKeyType> FromArrowArray<&DictionaryArray<K>> for DictArra
     }
 }
 
-fn nulls(nulls: Option<&NullBuffer>, nullable: bool) -> Validity {
+pub(crate) fn nulls(nulls: Option<&NullBuffer>, nullable: bool) -> Validity {
     if nullable {
         nulls
             .map(|nulls| {
@@ -909,8 +909,10 @@ mod tests {
         Arc::new(Date32Array::from(vec![18000_i32, 18001, 18002, 18003])),
     )]
     #[case::date64(
-        Arc::new(Date64Array::from(vec![Some(1555200000000), None, Some(1555286400000), Some(1555372800000)])),
-        Arc::new(Date64Array::from(vec![1555200000000_i64, 1555213600000, 1555286400000, 1555372800000])),
+        Arc::new(Date64Array::from(vec![Some(1555200000000), None, Some(1555286400000), Some(1555372800000)]
+        )),
+        Arc::new(Date64Array::from(vec![1555200000000_i64, 1555213600000, 1555286400000, 1555372800000]
+        )),
     )]
     fn test_temporal_array_conversion(
         #[case] nullable: Arc<dyn ArrowArray>,

vortex-array/src/arrow/executor/fixed_size_list.rs

@@ -12,8 +12,8 @@ use crate::ExecutionCtx;
 use crate::arrays::FixedSizeList;
 use crate::arrays::FixedSizeListArray;
 use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
-use crate::arrow::ArrowArrayExecutor;
 use crate::arrow::executor::validity::to_arrow_null_buffer;
+use crate::arrow::session::ArrowSessionExt;
 
 pub(super) fn to_arrow_fixed_list(
     array: ArrayRef,
@@ -44,10 +44,11 @@ fn list_to_list(
         list_size
     );
 
-    let elements = array
-        .elements()
-        .clone()
-        .execute_arrow(Some(elements_field.data_type()), ctx)?;
+    let elements = ctx.session().clone().arrow().execute_arrow(
+        array.elements().clone(),
+        Some(elements_field.as_ref()),
+        ctx,
+    )?;
     vortex_ensure!(
         elements_field.is_nullable() || elements.null_count() == 0,
         "Cannot convert FixedSizeListArray to non-nullable Arrow array when elements are nullable"

vortex-array/src/arrow/executor/list.rs

@@ -26,8 +26,8 @@ use crate::arrays::list::ListArrayExt;
 use crate::arrays::listview::ListViewArrayExt;
 use crate::arrays::listview::ListViewDataParts;
 use crate::arrays::listview::ListViewRebuildMode;
-use crate::arrow::ArrowArrayExecutor;
 use crate::arrow::executor::validity::to_arrow_null_buffer;
+use crate::arrow::session::ArrowSessionExt;
 use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
 use crate::dtype::NativePType;
@@ -96,10 +96,11 @@ fn list_to_list<O: OffsetSizeTrait + NativePType>(
         .to_buffer::<O>()
         .into_arrow_offset_buffer();
 
-    let elements = array
-        .elements()
-        .clone()
-        .execute_arrow(Some(elements_field.data_type()), ctx)?;
+    let elements = ctx.session().clone().arrow().execute_arrow(
+        array.elements().clone(),
+        Some(elements_field.as_ref()),
+        ctx,
+    )?;
     vortex_ensure!(
         elements_field.is_nullable() || elements.null_count() == 0,
         "Cannot convert to non-nullable Arrow array with null elements"
@@ -124,10 +125,11 @@ fn list_view_zctl<O: OffsetSizeTrait + NativePType>(
     assert!(array.is_zero_copy_to_list());
 
     if array.is_empty() {
-        let elements = array
-            .elements()
-            .clone()
-            .execute_arrow(Some(elements_field.data_type()), ctx)?;
+        let elements = ctx.session().clone().arrow().execute_arrow(
+            array.elements().clone(),
+            Some(elements_field.as_ref()),
+            ctx,
+        )?;
         return Ok(Arc::new(GenericListArray::<O>::new(
             Arc::clone(elements_field),
             OffsetBuffer::new_empty(),
@@ -176,7 +178,11 @@ fn list_view_zctl<O: OffsetSizeTrait + NativePType>(
     });
 
     // Extract the elements array.
-    let elements = elements.execute_arrow(Some(elements_field.data_type()), ctx)?;
+    let elements = ctx.session().clone().arrow().execute_arrow(
+        elements,
+        Some(elements_field.as_ref()),
+        ctx,
+    )?;
     vortex_ensure!(
         elements_field.is_nullable() || elements.null_count() == 0,
         "Cannot convert to non-nullable Arrow array with null elements"

vortex-array/src/arrow/executor/list_view.rs

@@ -15,8 +15,8 @@ use crate::arrays::ListView;
 use crate::arrays::ListViewArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::listview::ListViewDataParts;
-use crate::arrow::ArrowArrayExecutor;
 use crate::arrow::executor::validity::to_arrow_null_buffer;
+use crate::arrow::session::ArrowSessionExt;
 use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
 use crate::dtype::IntegerPType;
@@ -51,7 +51,11 @@ fn list_view_to_list_view<O: OffsetSizeTrait + IntegerPType>(
         ..
     } = array.into_data_parts();
 
-    let elements = elements.execute_arrow(Some(elements_field.data_type()), ctx)?;
+    let elements = ctx.session().clone().arrow().execute_arrow(
+        elements,
+        Some(elements_field.as_ref()),
+        ctx,
+    )?;
     vortex_ensure!(
         elements_field.is_nullable() || elements.null_count() == 0,
         "Elements field is non-nullable but elements array contains nulls"