Variant DType

encodings/sparse/src/canonical.rs

@@ -50,6 +50,7 @@ use vortex_buffer::buffer_mut;
 use vortex_error::VortexError;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
 use vortex_error::vortex_panic;
 
 use crate::ConstantArray;
@@ -117,6 +118,7 @@ pub(super) fn execute_sparse(
             execute_sparse_fixed_size_list(array, *nullability, ctx)?
         }
         DType::Extension(_ext_dtype) => todo!(),
+        DType::Variant(_) => vortex_bail!("Sparse canonicalization does not support Variant"),
     })
 }
 

fuzz/src/array/compare.rs

@@ -141,7 +141,7 @@ pub fn compare_canonical_array(
             )
             .into_array()
         }
-        d @ (DType::Null | DType::Extension(_)) => {
+        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/filter.rs

@@ -115,7 +115,7 @@ pub fn filter_canonical_array(array: &ArrayRef, filter: &[bool]) -> VortexResult
             }
             take_canonical_array_non_nullable_indices(array, indices.as_slice())
         }
-        d @ (DType::Null | DType::Extension(_)) => {
+        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/mod.rs

@@ -488,6 +488,8 @@ fn actions_for_dtype(dtype: &DType) -> HashSet<ActionType> {
             // Extension types delegate to storage dtype, support most operations
             ActionType::iter().collect()
         }
+        // Currently, no support at all
+        DType::Variant(_) => unreachable!("Variant dtype shouldn't be fuzzed"),
     }
 }
 

fuzz/src/array/search_sorted.rs

@@ -131,7 +131,7 @@ pub fn search_sorted_canonical_array(
                 .collect::<VortexResult<Vec<_>>>()?;
             scalar_vals.search_sorted(&scalar.cast(array.dtype())?, side)
         }
-        d @ (DType::Null | DType::Extension(_)) => {
+        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/slice.rs

@@ -114,7 +114,7 @@ pub fn slice_canonical_array(
                 .into_array(),
             )
         }
-        d @ (DType::Null | DType::Extension(_)) => {
+        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/sort.rs

@@ -81,7 +81,7 @@ pub fn sort_canonical_array(array: &ArrayRef) -> VortexResult<ArrayRef> {
             });
             take_canonical_array_non_nullable_indices(array, &sort_indices)
         }
-        d @ (DType::Null | DType::Extension(_)) => {
+        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

fuzz/src/array/take.rs

@@ -138,7 +138,7 @@ pub fn take_canonical_array(array: &ArrayRef, indices: &[Option<usize>]) -> Vort
             }
             Ok(builder.finish())
         }
-        d @ (DType::Null | DType::Extension(_)) => {
+        d @ (DType::Null | DType::Extension(_) | DType::Variant(_)) => {
             unreachable!("DType {d} not supported for fuzzing")
         }
     }

vortex-array/public-api.lock

@@ -10402,6 +10402,8 @@ pub vortex_array::dtype::DType::Struct(vortex_array::dtype::StructFields, vortex
 
 pub vortex_array::dtype::DType::Utf8(vortex_array::dtype::Nullability)
 
+pub vortex_array::dtype::DType::Variant(vortex_array::dtype::Nullability)
+
 impl vortex_array::dtype::DType
 
 pub const vortex_array::dtype::DType::BYTES: Self
@@ -10478,6 +10480,8 @@ pub fn vortex_array::dtype::DType::is_unsigned_int(&self) -> bool
 
 pub fn vortex_array::dtype::DType::is_utf8(&self) -> bool
 
+pub fn vortex_array::dtype::DType::is_variant(&self) -> bool
+
 pub fn vortex_array::dtype::DType::list(dtype: impl core::convert::Into<vortex_array::dtype::DType>, nullability: vortex_array::dtype::Nullability) -> Self
 
 pub fn vortex_array::dtype::DType::nullability(&self) -> vortex_array::dtype::Nullability
@@ -15330,6 +15334,8 @@ pub vortex_array::scalar::ScalarValue::Primitive(vortex_array::scalar::PValue)
 
 pub vortex_array::scalar::ScalarValue::Utf8(vortex_buffer::string::BufferString)
 
+pub vortex_array::scalar::ScalarValue::Variant(alloc::boxed::Box<vortex_array::scalar::Scalar>)
+
 impl vortex_array::scalar::ScalarValue
 
 pub fn vortex_array::scalar::ScalarValue::as_binary(&self) -> &vortex_buffer::ByteBuffer
@@ -15344,6 +15350,8 @@ pub fn vortex_array::scalar::ScalarValue::as_primitive(&self) -> &vortex_array::
 
 pub fn vortex_array::scalar::ScalarValue::as_utf8(&self) -> &vortex_buffer::string::BufferString
 
+pub fn vortex_array::scalar::ScalarValue::as_variant(&self) -> &vortex_array::scalar::Scalar
+
 pub fn vortex_array::scalar::ScalarValue::into_binary(self) -> vortex_buffer::ByteBuffer
 
 pub fn vortex_array::scalar::ScalarValue::into_bool(self) -> bool
@@ -15356,6 +15364,8 @@ pub fn vortex_array::scalar::ScalarValue::into_primitive(self) -> vortex_array::
 
 pub fn vortex_array::scalar::ScalarValue::into_utf8(self) -> vortex_buffer::string::BufferString
 
+pub fn vortex_array::scalar::ScalarValue::into_variant(self) -> vortex_array::scalar::Scalar
+
 impl vortex_array::scalar::ScalarValue
 
 pub fn vortex_array::scalar::ScalarValue::from_proto(value: &vortex_proto::scalar::ScalarValue, dtype: &vortex_array::dtype::DType, session: &vortex_session::VortexSession) -> vortex_error::VortexResult<core::option::Option<Self>>
@@ -15980,6 +15990,8 @@ pub fn vortex_array::scalar::Scalar::is_null(&self) -> bool
 
 pub fn vortex_array::scalar::Scalar::is_valid(&self) -> bool
 
+pub fn vortex_array::scalar::Scalar::is_variant_null(&self) -> core::option::Option<bool>
+
 pub fn vortex_array::scalar::Scalar::is_zero(&self) -> core::option::Option<bool>
 
 pub unsafe fn vortex_array::scalar::Scalar::new_unchecked(dtype: vortex_array::dtype::DType, value: core::option::Option<vortex_array::scalar::ScalarValue>) -> Self
@@ -16030,6 +16042,10 @@ pub fn vortex_array::scalar::Scalar::as_utf8(&self) -> vortex_array::scalar::Utf
 
 pub fn vortex_array::scalar::Scalar::as_utf8_opt(&self) -> core::option::Option<vortex_array::scalar::Utf8Scalar<'_>>
 
+pub fn vortex_array::scalar::Scalar::as_variant(&self) -> vortex_array::scalar::VariantScalar<'_>
+
+pub fn vortex_array::scalar::Scalar::as_variant_opt(&self) -> core::option::Option<vortex_array::scalar::VariantScalar<'_>>
+
 impl vortex_array::scalar::Scalar
 
 pub fn vortex_array::scalar::Scalar::binary(buffer: impl core::convert::Into<vortex_buffer::ByteBuffer>, nullability: vortex_array::dtype::Nullability) -> Self
@@ -16056,6 +16072,8 @@ pub fn vortex_array::scalar::Scalar::try_utf8<B>(str: B, nullability: vortex_arr
 
 pub fn vortex_array::scalar::Scalar::utf8<B>(str: B, nullability: vortex_array::dtype::Nullability) -> Self where B: core::convert::Into<vortex_buffer::string::BufferString>
 
+pub fn vortex_array::scalar::Scalar::variant(value: vortex_array::scalar::Scalar) -> Self
+
 impl vortex_array::scalar::Scalar
 
 pub fn vortex_array::scalar::Scalar::cast(&self, target_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
@@ -16720,6 +16738,32 @@ impl<'a> core::hash::Hash for vortex_array::scalar::Utf8Scalar<'a>
 
 pub fn vortex_array::scalar::Utf8Scalar<'a>::hash<__H: core::hash::Hasher>(&self, state: &mut __H)
 
+pub struct vortex_array::scalar::VariantScalar<'a>
+
+impl<'a> vortex_array::scalar::VariantScalar<'a>
+
+pub fn vortex_array::scalar::VariantScalar<'a>::dtype(&self) -> &'a vortex_array::dtype::DType
+
+pub fn vortex_array::scalar::VariantScalar<'a>::is_null(&self) -> bool
+
+pub fn vortex_array::scalar::VariantScalar<'a>::is_variant_null(&self) -> core::option::Option<bool>
+
+pub fn vortex_array::scalar::VariantScalar<'a>::value(&self) -> core::option::Option<&'a vortex_array::scalar::Scalar>
+
+impl core::fmt::Display for vortex_array::scalar::VariantScalar<'_>
+
+pub fn vortex_array::scalar::VariantScalar<'_>::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl<'a> core::clone::Clone for vortex_array::scalar::VariantScalar<'a>
+
+pub fn vortex_array::scalar::VariantScalar<'a>::clone(&self) -> vortex_array::scalar::VariantScalar<'a>
+
+impl<'a> core::fmt::Debug for vortex_array::scalar::VariantScalar<'a>
+
+pub fn vortex_array::scalar::VariantScalar<'a>::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl<'a> core::marker::Copy for vortex_array::scalar::VariantScalar<'a>
+
 pub trait vortex_array::scalar::ScalarTruncation: core::marker::Send + core::marker::Sized
 
 pub fn vortex_array::scalar::ScalarTruncation::from_scalar(value: vortex_array::scalar::Scalar) -> vortex_error::VortexResult<core::option::Option<Self>>

vortex-array/src/arrays/arbitrary.rs

@@ -162,7 +162,10 @@ fn random_array_chunk(
             random_fixed_size_list(u, elem_dtype, *list_size, *null, chunk_len)
         }
         DType::Extension(..) => {
-            todo!("Extension arrays are not implemented")
+            unimplemented!("Extension arrays are not implemented")
+        }
+        DType::Variant(_) => {
+            unimplemented!("Variant arrays are not implemented")
         }
     }
 }

vortex-array/src/arrays/constant/vtable/canonical.rs

@@ -164,6 +164,11 @@ pub(crate) fn constant_canonicalize(array: &ConstantArray) -> VortexResult<Canon
             let storage_self = ConstantArray::new(storage_scalar, array.len()).into_array();
             Canonical::Extension(ExtensionArray::new(ext_dtype.clone(), storage_self))
         }
+        DType::Variant(_) => {
+            unimplemented!(
+                "TODO(variant): canonicalization will use the child-array design in a follow-up"
+            )
+        }
     })
 }
 

vortex-array/src/builders/mod.rs

@@ -280,5 +280,8 @@ pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box<dyn ArrayBui
         DType::Extension(ext_dtype) => {
             Box::new(ExtensionBuilder::with_capacity(ext_dtype.clone(), capacity))
         }
+        DType::Variant(_) => {
+            unimplemented!()
+        }
     }
 }

vortex-array/src/builders/tests.rs

@@ -630,6 +630,7 @@ fn create_test_scalars_for_dtype(dtype: &DType, count: usize) -> Vec<Scalar> {
                 };
                 Scalar::extension_ref(ext_dtype.clone(), storage_scalar)
             }
+            DType::Variant(_) => continue,
         };
         scalars.push(scalar);
     }

vortex-array/src/canonical.rs

@@ -221,6 +221,9 @@ impl Canonical {
                 ext_dtype.clone(),
                 Canonical::empty(ext_dtype.storage_dtype()).into_array(),
             )),
+            DType::Variant(_) => {
+                vortex_panic!(InvalidArgument: "Canonical empty is not supported for Variant")
+            }
         }
     }
 

vortex-array/src/compute/conformance/consistency.rs

@@ -1235,6 +1235,7 @@ fn test_cast_slice_consistency(array: &ArrayRef) {
             )]
         }
         DType::Extension(_) => vec![], // Extension types typically only cast to themselves
+        DType::Variant(_) => unimplemented!(),
     };
 
     // Test each target dtype

vortex-array/src/dtype/arrow.rs

@@ -210,6 +210,14 @@ impl FromArrowType<(&DataType, Nullability)> for DType {
 
 impl FromArrowType<&Field> for DType {
     fn from_arrow(field: &Field) -> Self {
+        if field
+            .metadata()
+            .get("ARROW:extension:name")
+            .map(|s| s.as_str())
+            == Some("arrow.parquet.variant")
+        {
+            return DType::Variant(field.is_nullable().into());
+        }
         Self::from_arrow((field.data_type(), field.is_nullable().into()))
     }
 }
@@ -227,11 +235,23 @@ impl DType {
 
         let mut builder = SchemaBuilder::with_capacity(struct_dtype.names().len());
         for (field_name, field_dtype) in struct_dtype.names().iter().zip(struct_dtype.fields()) {
-            builder.push(FieldRef::from(Field::new(
-                field_name.as_ref(),
-                field_dtype.to_arrow_dtype()?,
-                field_dtype.is_nullable(),
-            )));
+            let field = if field_dtype.is_variant() {
+                let storage = DataType::Struct(variant_storage_fields_minimal());
+                Field::new(field_name.as_ref(), storage, field_dtype.is_nullable()).with_metadata(
+                    [(
+                        "ARROW:extension:name".to_owned(),
+                        "arrow.parquet.variant".to_owned(),
+                    )]
+                    .into(),
+                )
+            } else {
+                Field::new(
+                    field_name.as_ref(),
+                    field_dtype.to_arrow_dtype()?,
+                    field_dtype.is_nullable(),
+                )
+            };
+            builder.push(field);
         }
 
         Ok(builder.finish())
@@ -300,6 +320,9 @@ impl DType {
 
                 DataType::Struct(Fields::from(fields))
             }
+            DType::Variant(_) => vortex_bail!(
+                "DType::Variant requires Arrow Field metadata; use to_arrow_schema or a Field helper"
+            ),
             DType::Extension(ext_dtype) => {
                 // Try and match against the known extension DTypes.
                 if let Some(temporal) = ext_dtype.metadata_opt::<AnyTemporal>() {
@@ -332,6 +355,13 @@ impl DType {
     }
 }
 
+fn variant_storage_fields_minimal() -> Fields {
+    Fields::from(vec![
+        Field::new("metadata", DataType::Binary, false),
+        Field::new("value", DataType::Binary, true),
+    ])
+}
+
 #[cfg(test)]
 mod test {
     use arrow_schema::DataType;
@@ -399,6 +429,15 @@ mod test {
         );
     }
 
+    #[test]
+    fn test_variant_dtype_to_arrow_dtype_errors() {
+        let err = DType::Variant(Nullability::NonNullable)
+            .to_arrow_dtype()
+            .unwrap_err()
+            .to_string();
+        assert!(err.contains("Variant"));
+    }
+
     #[test]
     fn infer_nullable_list_element() {
         let list_non_nullable = DType::List(
@@ -455,6 +494,25 @@ mod test {
         );
     }
 
+    #[test]
+    fn test_schema_variant_field_metadata() {
+        let dtype = DType::struct_(
+            [("v", DType::Variant(Nullability::NonNullable))],
+            Nullability::NonNullable,
+        );
+        let schema = dtype.to_arrow_schema().unwrap();
+        let field = schema.field(0);
+        assert_eq!(
+            field
+                .metadata()
+                .get("ARROW:extension:name")
+                .map(|s| s.as_str()),
+            Some("arrow.parquet.variant")
+        );
+        assert!(matches!(field.data_type(), DataType::Struct(_)));
+        assert!(!field.is_nullable());
+    }
+
     #[rstest]
     #[should_panic]
     fn test_schema_conversion_panics(the_struct: StructFields) {