fixes

Signed-off-by: Baris Palaska <barispalaska@gmail.com>

Author: palaska

vortex-array/public-api.lock

@@ -8888,6 +8888,8 @@ pub use vortex_array::dtype::half
 
 pub mod vortex_array::dtype::arrow
 
+pub const vortex_array::dtype::arrow::ARROW_EXT_NAME_VARIANT: &str
+
 pub trait vortex_array::dtype::arrow::FromArrowType<T>: core::marker::Sized
 
 pub fn vortex_array::dtype::arrow::FromArrowType::from_arrow(value: T) -> Self

vortex-array/src/arrow/convert.rs

@@ -589,8 +589,7 @@ impl FromArrowArray<&dyn ArrowArray> for ArrayRef {
             DataType::Dictionary(key_type, _) => {
                 Ok(dict_from_arrow_with_session(array, key_type, nullable, session)?.into_array())
             }
-            // Other arrays don't carry child Fields, so the session-aware path is
-            // equivalent to the legacy one.
+            // Leaves: no child Fields to thread session through.
             _ => Self::from_arrow(array, nullable),
         }
     }
@@ -744,8 +743,8 @@ impl FromArrowArray<&RecordBatch> for ArrayRef {
     }
 }
 
-/// Inverse of `field_from_dtype` (in `dtype/arrow.rs`): rewrap `storage` as `ExtensionArray`
-/// when `field` carries `ARROW:extension:name` metadata for a registered extension.
+/// Rewrap `storage` as `ExtensionArray` if `field` carries `ARROW:extension:name`
+/// for a registered extension. Inverse of `field_from_dtype` in `dtype/arrow.rs`.
 fn wrap_extension_if_field_has_metadata(
     storage: ArrayRef,
     field: &Field,

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

@@ -99,10 +99,9 @@ impl ArrowArrayExecutor for ArrayRef {
             None => preferred_arrow_type(&self)?,
         };
 
-        // Extensions with a native Arrow mapping (temporal) keep their wrapper so
-        // `to_arrow_temporal` can read the metadata. Other extensions carry identity in
-        // Field metadata, so dispatch on the storage array. Mirror the discriminator used
-        // by `field_from_dtype` / `native_arrow_dtype_for_extension` in `dtype/arrow.rs`.
+        // Temporal extensions keep their wrapper so `to_arrow_temporal` can read the
+        // metadata. Other extensions unwrap to storage; their identity rides on Field
+        // metadata.
         if let DType::Extension(ext) = self.dtype()
             && ext.metadata_opt::<AnyTemporal>().is_none()
         {

vortex-array/src/dtype/arrow.rs

@@ -53,7 +53,9 @@ use crate::extension::datetime::Time;
 use crate::extension::datetime::TimeUnit;
 use crate::extension::datetime::Timestamp;
 
-const ARROW_EXT_NAME_VARIANT: &str = "arrow.parquet.variant";
+/// Canonical Arrow extension name for Parquet Variant — handled as `DType::Variant` rather
+/// than going through the extension registry.
+pub const ARROW_EXT_NAME_VARIANT: &str = "arrow.parquet.variant";
 
 /// Trait for converting Arrow types to Vortex types.
 pub trait FromArrowType<T>: Sized {
@@ -251,8 +253,7 @@ impl FromArrowType<&Field> for DType {
     }
 }
 
-/// Convert an Arrow Field to a [`DType`] with `dtypes` already borrowed from the session,
-/// so the handle is acquired once per schema rather than once per field.
+/// Convert a Field to a [`DType`]. Takes `dtypes` borrowed once per schema (not per field).
 fn dtype_from_field(field: &Field, dtypes: &DTypeSession) -> DType {
     if field
         .extension_type_name()
@@ -268,10 +269,8 @@ fn dtype_from_field(field: &Field, dtypes: &DTypeSession) -> DType {
     }
 }
 
-/// Resolve the [`ExtDTypeRef`] for an Arrow Field whose `ARROW:extension:name` metadata names
-/// a registered Vortex extension. Returns `None` for unregistered extensions, malformed
-/// metadata, or fields with no extension name; callers fall back to the storage representation
-/// and `tracing::warn!` reports the anomaly.
+/// Resolve the [`ExtDTypeRef`] for a Field whose `ARROW:extension:name` names a registered
+/// Vortex extension. Returns `None` for missing/unregistered/malformed metadata.
 pub(crate) fn resolve_extension_dtype(
     field: &Field,
     dtypes: &DTypeSession,
@@ -377,10 +376,8 @@ impl DType {
 
     /// Returns the Arrow [`DataType`] that best corresponds to this Vortex [`DType`].
     ///
-    /// Extensions without a native Arrow mapping (e.g. user-registered extensions like
-    /// `Vector`) degrade to their storage `DataType`; extension identity only survives when
-    /// emitted onto an Arrow `Field` (see [`Self::to_arrow_schema`]). Callers that must
-    /// reject non-temporal extensions should match on `DType::Extension` themselves.
+    /// Extensions without a native Arrow mapping degrade to their storage `DataType`;
+    /// identity only survives via `Field` metadata (see [`Self::to_arrow_schema`]).
     pub fn to_arrow_dtype(&self) -> VortexResult<DataType> {
         arrow_dtype_from_dtype(self)
     }
@@ -460,8 +457,7 @@ fn arrow_dtype_from_dtype(dtype: &DType) -> VortexResult<DataType> {
     })
 }
 
-/// Build an Arrow [`Field`], attaching `ARROW:extension:name` and, when present,
-/// `ARROW:extension:metadata` for extensions and Variant that have no native Arrow mapping.
+/// Build a Field, attaching extension/Variant metadata when there's no native Arrow mapping.
 fn field_from_dtype(name: &str, dtype: &DType) -> VortexResult<Field> {
     if dtype.is_variant() {
         let storage = DataType::Struct(variant_storage_fields_minimal());

vortex-python/src/arrays/from_arrow.rs

@@ -17,12 +17,12 @@ use vortex::array::arrays::ChunkedArray;
 use vortex::array::arrays::ExtensionArray;
 use vortex::array::arrow::FromArrowArray;
 use vortex::dtype::DType;
+use vortex::dtype::arrow::ARROW_EXT_NAME_VARIANT;
 use vortex::dtype::arrow::FromArrowType;
 use vortex::dtype::extension::ExtId;
 use vortex::dtype::session::DTypeSessionExt;
 use vortex::error::VortexError;
 use vortex::error::VortexResult;
-use vortex::error::vortex_err;
 use vortex::session::VortexSession;
 
 use crate::SESSION;
@@ -37,8 +37,7 @@ use crate::error::PyVortexResult;
 
 /// Convert a Python `pyarrow` array (including `pa.ExtensionArray`) into a Vortex array.
 ///
-/// Arrow's C ABI strips extension identity from the array layer — it lives on `Field`
-/// metadata, and a leaf `pa.ExtensionArray` has no enclosing field. We recover it from the
+/// The Arrow C ABI strips extension identity from leaf arrays; we recover it from the
 /// Python object via `extension_name` and `__arrow_ext_serialize__`.
 pub trait FromPyArrowArray: Sized {
     /// Convert a Python `pyarrow` array to a Vortex array.
@@ -78,22 +77,27 @@ impl FromPyArrowArray for ArrayRef {
     }
 }
 
-/// `__arrow_ext_serialize__` returns raw bytes; we pass them straight to the plugin.
-/// Base64 is only the encoding used in the Arrow Field-metadata string channel — going
-/// directly Python → registry skips that hop.
+/// Raw bytes from `__arrow_ext_serialize__` — no base64 (that's only for the
+/// Arrow Field-metadata string channel). Variant short-circuits to `None` so it surfaces
+/// as `DType::Variant` via the storage path, mirroring `dtype/arrow.rs::dtype_from_field`.
 fn extract_extension_info(py_array: &Bound<'_, PyAny>) -> PyResult<Option<(String, Vec<u8>)>> {
     let py = py_array.py();
     let py_type = py_array.getattr(intern!(py, "type"))?;
     if !py_type.is_instance(extension_type_class(py)?)? {
         return Ok(None);
     }
     let ext_name: String = py_type.getattr(intern!(py, "extension_name"))?.extract()?;
+    if ext_name == ARROW_EXT_NAME_VARIANT {
+        return Ok(None);
+    }
     let ext_meta_bytes: Vec<u8> = py_type
         .call_method0(intern!(py, "__arrow_ext_serialize__"))?
         .extract()?;
     Ok(Some((ext_name, ext_meta_bytes)))
 }
 
+/// Soft fallback to storage on registry miss or malformed metadata, mirroring
+/// `dtype/arrow.rs::resolve_extension_dtype`.
 fn wrap_with_extension(
     storage: ArrayRef,
     ext_name: &str,
@@ -102,11 +106,20 @@ fn wrap_with_extension(
 ) -> VortexResult<ArrayRef> {
     let ext_id = ExtId::new(ext_name);
     let dtypes = session.dtypes();
-    let plugin = dtypes
-        .registry()
-        .find(&ext_id)
-        .ok_or_else(|| vortex_err!("extension `{ext_name}` is not registered on the session"))?;
-    let ext_dtype = plugin.deserialize(ext_meta_bytes, storage.dtype().clone())?;
+    let Some(plugin) = dtypes.registry().find(&ext_id) else {
+        log::warn!("pyarrow extension {ext_name:?} not registered on session; using storage dtype");
+        return Ok(storage);
+    };
+    let ext_dtype = match plugin.deserialize(ext_meta_bytes, storage.dtype().clone()) {
+        Ok(dt) => dt,
+        Err(e) => {
+            log::warn!(
+                "pyarrow extension {ext_name:?} failed to deserialize metadata ({e}); \
+                 using storage dtype",
+            );
+            return Ok(storage);
+        }
+    };
     Ok(ExtensionArray::try_new(ext_dtype, storage)?.into_array())
 }
 
@@ -135,28 +148,49 @@ pub(super) fn from_arrow(obj: &Borrowed<'_, '_, PyAny>) -> PyVortexResult<PyArra
         Ok(PyArrayRef::from(enc_array))
     } else if obj.is_instance(chunked_array)? {
         let chunks: Vec<Bound<PyAny>> = obj.getattr(intern!(py, "chunks"))?.extract()?;
+        // ChunkedArray has a uniform type — peek extension identity once and reuse.
+        let bound = obj.to_owned();
+        let ext_info = extract_extension_info(&bound)?;
         let encoded_chunks = chunks
             .iter()
-            .map(|a| {
-                let arrow_array = ArrowArrayData::from_pyarrow(&a.as_borrowed()).map(make_array)?;
-                ArrayRef::from_arrow(arrow_array.as_ref(), false).map_err(PyVortexError::from)
+            .map(|chunk| {
+                let arrow_array =
+                    ArrowArrayData::from_pyarrow(&chunk.as_borrowed()).map(make_array)?;
+                let storage = ArrayRef::from_arrow_with_session(
+                    arrow_array.as_ref(),
+                    arrow_array.is_nullable(),
+                    &SESSION,
+                )
+                .map_err(PyVortexError::from)?;
+                match &ext_info {
+                    None => Ok(storage),
+                    Some((name, meta)) => wrap_with_extension(storage, name, meta, &SESSION)
+                        .map_err(|e| PyVortexError::from(e).into()),
+                }
             })
-            .collect::<PyVortexResult<Vec<_>>>()?;
-        let dtype: DType = obj
-            .getattr(intern!(py, "type"))
-            .and_then(|v| DataType::from_pyarrow(&v.as_borrowed()))
-            .map(|dt| DType::from_arrow(&Field::new("_", dt, false)))?;
+            .collect::<PyResult<Vec<_>>>()?;
+        let dtype: DType = if let Some(first) = encoded_chunks.first() {
+            first.dtype().clone()
+        } else {
+            // Empty array: `obj.type` over the C ABI loses extension metadata, so we
+            // recover only the storage dtype.
+            obj.getattr(intern!(py, "type"))
+                .and_then(|v| DataType::from_pyarrow(&v.as_borrowed()))
+                .map(|dt| DType::from_arrow_with_session(&Field::new("_", dt, false), &SESSION))?
+        };
         Ok(PyArrayRef::from(
             ChunkedArray::try_new(encoded_chunks, dtype)?.into_array(),
         ))
     } else if obj.is_instance(table)? {
+        // The C ABI Stream carries Field metadata on the schema — session-aware
+        // conversion recovers extensions directly, no Python peek needed.
         let array_stream = ArrowArrayStreamReader::from_pyarrow(&obj.as_borrowed())?;
-        let dtype = DType::from_arrow(array_stream.schema());
+        let dtype = DType::from_arrow_with_session(array_stream.schema(), &SESSION);
         let chunks = array_stream
             .into_iter()
             .map(|b| {
                 b.map_err(VortexError::from)
-                    .and_then(|b| ArrayRef::from_arrow(b, false))
+                    .and_then(|b| ArrayRef::from_arrow_with_session(b, false, &SESSION))
             })
             .collect::<VortexResult<Vec<_>>>()?;
         Ok(PyArrayRef::from(

vortex-python/test/test_arrow_extension.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+"""Round-trip tests for `vx.array` against pyarrow inputs carrying Vortex
+extension identity over the Arrow C ABI."""
+
+from __future__ import annotations
+
+import base64
+
+import pyarrow as pa
+
+import vortex as vx
+
+# vortex.timestamp wire format: u8 unit_tag + u16 LE tz_len (us=1, no tz).
+# See vortex-array/src/extension/datetime/timestamp.rs::serialize_metadata.
+_TIMESTAMP_US_METADATA = bytes([1, 0, 0])
+
+
+class VortexTimestampType(pa.ExtensionType):
+    """A pyarrow `ExtensionType` matching Vortex's `vortex.timestamp` extension."""
+
+    def __init__(self, unit: str = "us"):
+        # pyarrow calls `__arrow_ext_serialize__` in __init__, so set `_unit` first.
+        self._unit = unit
+        pa.ExtensionType.__init__(self, pa.int64(), "vortex.timestamp")
+
+    def __arrow_ext_serialize__(self) -> bytes:
+        unit_tag = {"ns": 0, "us": 1, "ms": 2, "s": 3}[self._unit]
+        return bytes([unit_tag, 0, 0])
+
+    @classmethod
+    def __arrow_ext_deserialize__(cls, storage_type, serialized):  # noqa: ARG003
+        unit_tag = serialized[0]
+        unit = {0: "ns", 1: "us", 2: "ms", 3: "s"}[unit_tag]
+        return cls(unit)
+
+
+def test_chunked_extension_array_uses_session_for_leaf_extension_type():
+    ext_type = VortexTimestampType()
+    storage = pa.array([1, 2, 3], type=pa.int64())
+    arrow = pa.chunked_array([pa.ExtensionArray.from_storage(ext_type, storage)])
+    array = vx.array(arrow)
+    assert isinstance(array, vx.ChunkedArray)
+    assert repr(array.dtype) == repr(vx.timestamp("us"))
+    assert repr(array.chunks()[0].dtype) == repr(vx.timestamp("us"))
+
+
+def test_table_uses_session_for_extension_field_metadata():
+    field = pa.field("ts", pa.int64(), nullable=False).with_metadata(
+        {
+            b"ARROW:extension:name": b"vortex.timestamp",
+            b"ARROW:extension:metadata": base64.b64encode(_TIMESTAMP_US_METADATA),
+        }
+    )
+    table = pa.Table.from_arrays(
+        [pa.array([1, 2, 3], type=pa.int64())],
+        schema=pa.schema([field]),
+    )
+    array = vx.array(table)
+    expected = vx.struct({"ts": vx.timestamp("us")})
+    assert isinstance(array, vx.ChunkedArray)
+    assert repr(array.dtype) == repr(expected)
+    assert repr(array.chunks()[0].dtype) == repr(expected)