diff --git a/vortex-arrow/src/convert.rs b/vortex-arrow/src/convert.rs index 40d46025bd7..3fe301b6fe2 100644 --- a/vortex-arrow/src/convert.rs +++ b/vortex-arrow/src/convert.rs @@ -13,6 +13,7 @@ use arrow_array::GenericByteArray; use arrow_array::GenericByteViewArray; use arrow_array::GenericListArray; use arrow_array::GenericListViewArray; +use arrow_array::MapArray as ArrowMapArray; use arrow_array::NullArray as ArrowNullArray; use arrow_array::OffsetSizeTrait; use arrow_array::PrimitiveArray as ArrowPrimitiveArray; @@ -65,16 +66,22 @@ use vortex_array::arrays::DictArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::MapArray; use vortex_array::arrays::NullArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Struct; use vortex_array::arrays::StructArray; use vortex_array::arrays::TemporalArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::FieldNames; use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::MapDType; use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::i256; use vortex_array::extension::datetime::TimeUnit; @@ -477,6 +484,77 @@ impl FromArrowArray<&ArrowFixedSizeListArray> for ArrayRef { } } +pub(crate) fn map_from_arrow_parts( + entries: ArrayRef, + offsets: &OffsetBuffer, + arrow_nulls: Option<&NullBuffer>, + keys_sorted: bool, + nullable: bool, +) -> VortexResult { + let DType::Struct(struct_dtype, Nullability::NonNullable) = entries.dtype() else { + vortex_bail!( + "Arrow map entries must import as non-nullable struct, got {}", + entries.dtype() + ); + }; + vortex_ensure!( + struct_dtype.nfields() == 2, + "Arrow map entries struct must contain exactly two fields" + ); + + let key_dtype = struct_dtype + .field_by_index(0) + .ok_or_else(|| vortex_err!("Arrow map entries struct missing key field"))?; + vortex_ensure!( + !key_dtype.is_nullable(), + "Arrow map key field must be non-nullable" + ); + let value_dtype = struct_dtype + .field_by_index(1) + .ok_or_else(|| vortex_err!("Arrow map entries struct missing value field"))?; + let map_dtype = MapDType::try_new(key_dtype, value_dtype, keys_sorted)?; + let entries_struct = entries.as_::(); + let entries = StructArray::try_new( + FieldNames::from(["key", "value"]), + vec![ + entries_struct.unmasked_field(0).clone(), + entries_struct.unmasked_field(1).clone(), + ], + entries.len(), + Validity::NonNullable, + )? + .into_array(); + + let len = offsets.len() - 1; + let sizes = Buffer::::from_iter(offsets.windows(2).map(|window| window[1] - window[0])) + .into_array(); + let offsets = offsets.inner().slice(0, len).into_array(); + let validity = nulls(arrow_nulls, nullable)?; + let entries = ListViewArray::try_new(entries, offsets, sizes, validity)?; + // SAFETY: Arrow Map offsets are sorted list offsets. Vortex sizes are adjacent offset + // differences, so every view's end equals the next view's start. Sliced Arrow arrays may + // leave leading or trailing entries unreferenced, which zero-copy-to-list permits. + let entries = unsafe { entries.with_zero_copy_to_list(true) }; + + Ok(MapArray::try_new(map_dtype, entries)?.into_array()) +} + +impl FromArrowArray<&ArrowMapArray> for ArrayRef { + fn from_arrow(array: &ArrowMapArray, nullable: bool) -> VortexResult { + let DataType::Map(_, keys_sorted) = array.data_type() else { + vortex_panic!("Invalid data type for MapArray: {}", array.data_type()); + }; + let entries = Self::from_arrow(array.entries(), false)?; + map_from_arrow_parts( + entries, + array.offsets(), + array.nulls(), + *keys_sorted, + nullable, + ) + } +} + impl FromArrowArray<&ArrowNullArray> for ArrayRef { fn from_arrow(value: &ArrowNullArray, nullable: bool) -> VortexResult { vortex_ensure!( @@ -546,9 +624,7 @@ impl FromArrowArray<&dyn ArrowArray> for ArrayRef { DataType::ListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::LargeListView(_) => Self::from_arrow(array.as_list_view::(), nullable), DataType::FixedSizeList(..) => Self::from_arrow(array.as_fixed_size_list(), nullable), - DataType::Map(..) => { - vortex_bail!("Arrow MapArray conversion is not yet supported") - } + DataType::Map(..) => Self::from_arrow(array.as_map(), nullable), DataType::Null => Self::from_arrow(as_null_array(array), nullable), DataType::Timestamp(u, _) => match u { ArrowTimeUnit::Second => { diff --git a/vortex-arrow/src/executor/map.rs b/vortex-arrow/src/executor/map.rs new file mode 100644 index 00000000000..c6b71057a87 --- /dev/null +++ b/vortex-arrow/src/executor/map.rs @@ -0,0 +1,518 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use arrow_array::Array as ArrowArray; +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::MapArray as ArrowMapArray; +use arrow_array::cast::AsArray; +use arrow_schema::DataType; +use arrow_schema::FieldRef; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::Map; +use vortex_array::arrays::MapArray; +use vortex_array::arrays::map::MapArrayExt; +use vortex_array::arrays::map::MapArraySlotsExt; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::executor::list::to_arrow_list; + +/// Converts a Vortex Map array into an Arrow [`MapArray`](ArrowMapArray). +/// +/// The Map entries are exported through the ListView-to-Arrow List path and then repackaged with +/// Arrow's Map field metadata. When the requested Arrow field asserts sorted keys, the Vortex Map +/// dtype must already make the same assertion. +pub(super) fn to_arrow_map( + array: ArrayRef, + entries_field: &FieldRef, + keys_sorted: bool, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let array = match array.try_downcast::() { + Ok(map) => map, + Err(array) => array.execute::(ctx)?, + }; + + vortex_ensure!( + !keys_sorted || array.keys_sorted(), + "Cannot convert unsorted Vortex map to Arrow MapArray with keys_sorted=true" + ); + + let entries = array.entries().clone(); + let entries_list_type = DataType::List(Arc::clone(entries_field)); + let entries_list = to_arrow_list::(entries, entries_field, ctx)?; + vortex_ensure!( + entries_list.data_type() == &entries_list_type, + "Arrow Map entries converted to {}, expected {entries_list_type}", + entries_list.data_type() + ); + + let entries_list = entries_list.as_list::(); + let entries = entries_list.values().as_struct().clone(); + let map = ArrowMapArray::try_new( + Arc::clone(entries_field), + entries_list.offsets().clone(), + entries, + entries_list.nulls().cloned(), + keys_sorted, + )?; + + Ok(Arc::new(map)) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::Array as ArrowArray; + use arrow_array::ArrayRef as ArrowArrayRef; + use arrow_array::FixedSizeBinaryArray; + use arrow_array::Int32Array; + use arrow_array::MapArray as ArrowMapArray; + use arrow_array::StringArray; + use arrow_array::StructArray as ArrowStructArray; + use arrow_array::builder::Int32Builder; + use arrow_array::builder::MapBuilder as ArrowMapBuilder; + use arrow_array::builder::StringBuilder; + use arrow_array::cast::AsArray; + use arrow_buffer::NullBuffer; + use arrow_buffer::OffsetBuffer; + use arrow_buffer::ScalarBuffer; + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::Fields; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::ListView; + use vortex_array::arrays::Map as VortexMap; + use vortex_array::arrays::listview::ListViewArrayExt; + use vortex_array::arrays::map::MapArraySlotsExt; + use vortex_array::builders::ArrayBuilder; + use vortex_array::builders::MapBuilder; + use vortex_array::dtype::DType; + use vortex_array::dtype::MapDType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::Nullability::Nullable; + use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + use vortex_mask::Mask; + + use crate::FromArrowArray as _; + use crate::session::ArrowSessionExt as _; + + fn i32_utf8_map_field(keys_sorted: bool, nullable: bool) -> Field { + let fields = Fields::from(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, true), + ]); + Field::new( + "maps", + DataType::Map( + Arc::new(Field::new_struct("entries", fields, false)), + keys_sorted, + ), + nullable, + ) + } + + fn i32_utf8_map_array( + offsets: Vec, + keys: Vec, + values: Vec>, + nulls: Option, + keys_sorted: bool, + ) -> VortexResult { + let fields = Fields::from(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, true), + ]); + let entries = ArrowStructArray::try_new( + fields.clone(), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(StringArray::from(values)), + ], + None, + )?; + ArrowMapArray::try_new( + Arc::new(Field::new_struct("entries", fields, false)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + entries, + nulls, + keys_sorted, + ) + .map_err(Into::into) + } + + #[test] + fn map_roundtrip_preserves_null_empty_duplicate_and_sorted_rows() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + let field = i32_utf8_map_field(true, true); + let arrow = i32_utf8_map_array( + vec![0, 2, 2, 2, 4], + vec![1, 2, 1, 1], + vec![Some("one"), None, Some("dup-old"), Some("dup-new")], + Some(NullBuffer::from_iter([true, false, true, true])), + true, + )?; + + let vortex = session.from_arrow_array(Arc::new(arrow), &field)?; + assert_eq!( + vortex.dtype(), + &DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(Nullable), + true, + Nullable, + )? + ); + + let exported = session.execute_arrow(vortex, Some(&field), &mut ctx)?; + let map = exported.as_map(); + assert_eq!(map.value_offsets(), &[0, 2, 2, 2, 4]); + assert!(map.is_valid(0)); + assert!(map.is_null(1)); + assert!(map.is_valid(2)); + assert!(map.is_valid(3)); + assert_eq!( + map.keys() + .as_primitive::() + .values(), + &[1, 2, 1, 1] + ); + let values = map.values().as_string::(); + assert_eq!(values.value(0), "one"); + assert!(values.is_null(1)); + assert_eq!(values.value(2), "dup-old"); + assert_eq!(values.value(3), "dup-new"); + + Ok(()) + } + + #[test] + fn legacy_from_arrow_map_imports_canonical_map() -> VortexResult<()> { + let arrow = i32_utf8_map_array( + vec![0, 1, 2], + vec![1, 2], + vec![Some("one"), Some("two")], + None, + false, + )?; + let vortex = ArrayRef::from_arrow(&arrow, false)?; + + assert!(vortex.is::()); + assert_eq!( + vortex.dtype(), + &DType::map( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(Nullable), + false, + NonNullable, + )? + ); + + Ok(()) + } + + #[test] + fn arrow_rs_map_builder_default_field_names_import() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + let mut builder = ArrowMapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + + builder.keys().append_value("joe"); + builder.values().append_value(1); + builder.append(true)?; + builder.keys().append_value("blogs"); + builder.values().append_value(2); + builder.keys().append_value("foo"); + builder.values().append_value(4); + builder.append(true)?; + builder.append(true)?; + builder.append(false)?; + + let arrow = builder.finish(); + let DataType::Map(entries_field, keys_sorted) = arrow.data_type() else { + panic!("expected Arrow map dtype, got {}", arrow.data_type()); + }; + assert!(!*keys_sorted); + assert_eq!(entries_field.name(), "entries"); + let DataType::Struct(fields) = entries_field.data_type() else { + panic!("expected Arrow map entries struct, got {entries_field:?}"); + }; + assert_eq!(fields[0].name(), "keys"); + assert_eq!(fields[1].name(), "values"); + + let expected_dtype = DType::map( + DType::Utf8(NonNullable), + DType::Primitive(PType::I32, Nullable), + false, + Nullable, + )?; + let legacy = ArrayRef::from_arrow(&arrow, true)?; + assert_eq!(legacy.dtype(), &expected_dtype); + + let field = Field::new("maps", arrow.data_type().clone(), true); + let vortex = session.from_arrow_array(Arc::new(arrow), &field)?; + assert_eq!(vortex.dtype(), &expected_dtype); + let row0 = vortex.execute_scalar(0, &mut ctx)?; + let row0_entries = row0.as_map().entries().collect::>(); + assert_eq!(row0_entries.len(), 1); + assert_eq!(row0_entries[0].0, Scalar::utf8("joe", NonNullable)); + assert_eq!(row0_entries[0].1, Scalar::primitive(1_i32, Nullable)); + let row1 = vortex.execute_scalar(1, &mut ctx)?; + let row1_entries = row1.as_map().entries().collect::>(); + assert_eq!(row1_entries.len(), 2); + assert_eq!(row1_entries[0].0, Scalar::utf8("blogs", NonNullable)); + assert_eq!(row1_entries[0].1, Scalar::primitive(2_i32, Nullable)); + assert!(vortex.execute_scalar(2, &mut ctx)?.as_map().is_empty()); + assert!(vortex.execute_scalar(3, &mut ctx)?.is_null()); + + Ok(()) + } + + #[test] + fn arrow_map_import_sets_zero_copy_to_list() -> VortexResult<()> { + let vortex_session = array_session(); + let session = vortex_session.arrow(); + let field = i32_utf8_map_field(false, false); + let arrow = i32_utf8_map_array( + vec![0, 2, 3], + vec![1, 2, 3], + vec![Some("one"), Some("two"), Some("three")], + None, + false, + )?; + + let vortex = session.from_arrow_array(Arc::new(arrow), &field)?; + let map = vortex.as_::(); + assert!(map.entries().as_::().is_zero_copy_to_list()); + + let sliced = i32_utf8_map_array( + vec![0, 1, 3, 4], + vec![9, 10, 11, 12], + vec![Some("nine"), Some("ten"), Some("eleven"), Some("twelve")], + None, + false, + )?; + let sliced: ArrowArrayRef = Arc::new(sliced.slice(1, 2)); + + let vortex = session.from_arrow_array(sliced, &field)?; + let map = vortex.as_::(); + let entries = map.entries().as_::(); + assert!(entries.is_zero_copy_to_list()); + assert_eq!(entries.offset_at(0), 1); + assert_eq!(entries.offset_at(1), 3); + + Ok(()) + } + + #[test] + fn sliced_arrow_map_import_preserves_nonzero_offsets() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + let field = i32_utf8_map_field(true, false); + let arrow = i32_utf8_map_array( + vec![0, 1, 3, 4], + vec![9, 10, 11, 12], + vec![Some("nine"), Some("ten"), Some("eleven"), Some("twelve")], + None, + true, + )?; + let sliced: ArrowArrayRef = Arc::new(arrow.slice(1, 2)); + + let vortex = session.from_arrow_array(sliced, &field)?; + let map = vortex.as_::(); + let entries = map.entries().as_::(); + assert_eq!(entries.offset_at(0), 1); + assert_eq!(entries.offset_at(1), 3); + assert_eq!(entries.size_at(0), 2); + assert_eq!(entries.size_at(1), 1); + + let exported = session.execute_arrow(vortex, Some(&field), &mut ctx)?; + let map = exported.as_map(); + + assert_eq!(map.value_offsets(), &[1, 3, 4]); + assert_eq!( + map.value(0) + .column(0) + .as_primitive::() + .values(), + &[10, 11] + ); + assert_eq!( + map.value(1) + .column(0) + .as_primitive::() + .values(), + &[12] + ); + + Ok(()) + } + + #[test] + fn filtered_map_with_entry_gap_exports_correct_rows() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + let map_dtype = MapDType::try_new( + DType::Primitive(PType::I32, NonNullable), + DType::Utf8(Nullable), + false, + )?; + let dtype = DType::Map(map_dtype.clone(), Nullable); + let mut builder = MapBuilder::::with_capacity(map_dtype, Nullable, 3); + let rows: [&[(i32, &str)]; 3] = [ + &[(1, "a"), (2, "b")], + &[(3, "c"), (4, "d"), (5, "e")], + &[(6, "f")], + ]; + for row in rows { + let entries = row + .iter() + .map(|(key, value)| { + ( + Scalar::primitive(*key, NonNullable), + Scalar::utf8(*value, Nullable), + ) + }) + .collect::>(); + builder.append_scalar(&Scalar::try_map(dtype.clone(), entries)?)?; + } + let source = builder.finish_into_map().into_array(); + + // Dropping the non-empty middle row leaves a gap in the entry elements, so the export + // must rebuild the entries list instead of reusing the builder's contiguous layout. + let filtered = source.filter(Mask::from_iter([true, false, true]))?; + + let field = i32_utf8_map_field(false, true); + let exported = session.execute_arrow(filtered, Some(&field), &mut ctx)?; + let map = exported.as_map(); + + assert_eq!(map.value_offsets(), &[0, 2, 3]); + assert_eq!( + map.keys() + .as_primitive::() + .values(), + &[1, 2, 6] + ); + let values = map.values().as_string::(); + assert_eq!(values.value(0), "a"); + assert_eq!(values.value(1), "b"); + assert_eq!(values.value(2), "f"); + + Ok(()) + } + + #[test] + fn unsorted_map_rejects_sorted_arrow_target() -> VortexResult<()> { + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + let source_field = i32_utf8_map_field(false, false); + let target_field = i32_utf8_map_field(true, false); + let arrow = i32_utf8_map_array( + vec![0, 2], + vec![2, 1], + vec![Some("two"), Some("one")], + None, + false, + )?; + let vortex = session.from_arrow_array(Arc::new(arrow), &source_field)?; + + let error = session + .execute_arrow(vortex, Some(&target_field), &mut ctx) + .unwrap_err(); + assert!(error.to_string().contains("keys_sorted=true")); + + Ok(()) + } + + #[test] + fn malformed_map_field_errors_without_panic() -> VortexResult<()> { + let session = array_session(); + let arrow_session = session.arrow(); + let arrow = i32_utf8_map_array(vec![0, 1], vec![1], vec![Some("one")], None, false)?; + let fields = Fields::from(vec![ + Field::new("key", DataType::Int32, true), + Field::new("value", DataType::Utf8, true), + ]); + let bad_field = Field::new( + "maps", + DataType::Map(Arc::new(Field::new_struct("entries", fields, false)), false), + false, + ); + + let error = arrow_session + .from_arrow_array(Arc::new(arrow), &bad_field) + .unwrap_err(); + assert!(error.to_string().contains("key field must be non-nullable")); + + Ok(()) + } + + #[test] + fn map_roundtrip_preserves_nested_uuid_fields() -> VortexResult<()> { + use arrow_schema::extension::Uuid as ArrowUuid; + + let vortex_session = array_session(); + let mut ctx = vortex_session.create_execution_ctx(); + let session = vortex_session.arrow(); + + let mut key_field = Field::new("key", DataType::FixedSizeBinary(16), false); + key_field.try_with_extension_type(ArrowUuid)?; + let mut value_field = Field::new("value", DataType::FixedSizeBinary(16), true); + value_field.try_with_extension_type(ArrowUuid)?; + let fields = Fields::from(vec![key_field, value_field]); + let entries_field = Arc::new(Field::new_struct("entries", fields.clone(), false)); + let field = Field::new( + "ids", + DataType::Map(Arc::clone(&entries_field), true), + false, + ); + let keys = FixedSizeBinaryArray::try_from_iter( + [ + b"0123456789abcdef".as_slice(), + b"fedcba9876543210".as_slice(), + ] + .into_iter(), + )?; + let values = FixedSizeBinaryArray::try_from_sparse_iter_with_size( + [Some(b"aaaaaaaaaaaaaaaa".as_slice()), None].into_iter(), + 16, + )?; + let entries = + ArrowStructArray::try_new(fields, vec![Arc::new(keys), Arc::new(values)], None)?; + let arrow = ArrowMapArray::try_new( + entries_field, + OffsetBuffer::new(ScalarBuffer::from(vec![0, 2])), + entries, + None, + true, + )?; + + let vortex = session.from_arrow_array(Arc::new(arrow), &field)?; + let DType::Map(map_dtype, NonNullable) = vortex.dtype() else { + panic!("expected map dtype, got {}", vortex.dtype()); + }; + assert!(map_dtype.key_dtype().is_extension()); + assert!(map_dtype.value_dtype().is_extension()); + + let exported = session.execute_arrow(vortex, Some(&field), &mut ctx)?; + assert_eq!(exported.data_type(), field.data_type()); + assert_eq!(exported.as_map().value_offsets(), &[0, 2]); + + Ok(()) + } +} diff --git a/vortex-arrow/src/executor/mod.rs b/vortex-arrow/src/executor/mod.rs index cae67f3143b..aacc93b869a 100644 --- a/vortex-arrow/src/executor/mod.rs +++ b/vortex-arrow/src/executor/mod.rs @@ -14,6 +14,7 @@ mod dictionary; mod fixed_size_list; mod list; mod list_view; +mod map; pub mod null; pub mod primitive; mod run_end; @@ -51,6 +52,7 @@ use crate::executor::dictionary::to_arrow_dictionary; use crate::executor::fixed_size_list::to_arrow_fixed_list; use crate::executor::list::to_arrow_list; use crate::executor::list_view::to_arrow_list_view; +use crate::executor::map::to_arrow_map; use crate::executor::null::to_arrow_null; use crate::executor::primitive::to_arrow_primitive; use crate::executor::run_end::to_arrow_run_end; @@ -191,7 +193,9 @@ pub(crate) fn execute_arrow_naive( dt @ (DataType::Date32 | DataType::Date64) => to_arrow_date(array, dt, ctx), dt @ (DataType::Time32(_) | DataType::Time64(_)) => to_arrow_time(array, dt, ctx), dt @ DataType::Timestamp(..) => to_arrow_timestamp(array, dt, ctx), - DataType::Map(..) => vortex_bail!("Arrow MapArray conversion is not yet supported"), + DataType::Map(entries_field, keys_sorted) => { + to_arrow_map(array, entries_field, *keys_sorted, ctx) + } DataType::FixedSizeBinary(_) | DataType::Duration(_) | DataType::Interval(_) diff --git a/vortex-arrow/src/scalar.rs b/vortex-arrow/src/scalar.rs index 1b8f2c01ed1..09fe7cfcdee 100644 --- a/vortex-arrow/src/scalar.rs +++ b/vortex-arrow/src/scalar.rs @@ -13,13 +13,13 @@ use arrow_schema::Field; use arrow_schema::Fields; use vortex_array::dtype::DType; use vortex_array::dtype::PType; +use vortex_array::dtype::i256; use vortex_array::extension::datetime::AnyTemporal; use vortex_array::extension::datetime::TemporalMetadata; use vortex_array::extension::datetime::TimeUnit; use vortex_array::scalar::BinaryScalar; use vortex_array::scalar::BoolScalar; use vortex_array::scalar::DecimalScalar; -use vortex_array::scalar::DecimalValue; use vortex_array::scalar::ExtScalar; use vortex_array::scalar::MapScalar; use vortex_array::scalar::PrimitiveScalar; @@ -75,12 +75,14 @@ impl ToArrowDatum for Scalar { DType::Decimal(..) => decimal_to_arrow(value.as_decimal()), DType::Utf8(_) => utf8_to_arrow(value.as_utf8()), DType::Binary(_) => binary_to_arrow(value.as_binary()), - DType::List(..) => unimplemented!("list scalar conversion"), - DType::FixedSizeList(..) => unimplemented!("fixed-size list scalar conversion"), + DType::List(..) => vortex_bail!("list scalar conversion is not supported"), + DType::FixedSizeList(..) => { + vortex_bail!("fixed-size list scalar conversion is not supported") + } DType::Map(..) => map_to_arrow(value.as_map()), - DType::Struct(..) => unimplemented!("struct scalar conversion"), - DType::Union(..) => unimplemented!("union scalar conversion"), - DType::Variant(_) => unimplemented!("Variant scalar conversion"), + DType::Struct(..) => vortex_bail!("struct scalar conversion is not supported"), + DType::Union(..) => vortex_bail!("union scalar conversion is not supported"), + DType::Variant(_) => vortex_bail!("Variant scalar conversion is not supported"), DType::Extension(..) => extension_to_arrow(value.as_extension()), } } @@ -110,20 +112,50 @@ fn primitive_to_arrow(scalar: PrimitiveScalar<'_>) -> Result, Vor /// Convert a [`DecimalScalar`] to an Arrow [`Datum`]. fn decimal_to_arrow(scalar: DecimalScalar<'_>) -> Result, VortexError> { + let DType::Decimal(decimal_dtype, _) = scalar.dtype() else { + vortex_bail!("Expected decimal scalar, got {}", scalar.dtype()); + }; + let precision = decimal_dtype.precision(); + let scale = decimal_dtype.scale(); // TODO(joe): Replace with decimal32, etc. once Arrow supports them. match scalar.decimal_value() { - Some(DecimalValue::I8(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))), - Some(DecimalValue::I16(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))), - Some(DecimalValue::I32(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))), - Some(DecimalValue::I64(v)) => Ok(Arc::new(Decimal128Array::new_scalar(v as i128))), - Some(DecimalValue::I128(v128)) => Ok(Arc::new(Decimal128Array::new_scalar(v128))), - Some(DecimalValue::I256(v256)) => Ok(Arc::new(Decimal256Array::new_scalar(v256.into()))), - None => Ok(Arc::new(arrow_array::Scalar::new( - Decimal128Array::new_null(SCALAR_ARRAY_LEN), - ))), + Some(value) => { + let value = value.as_i256(); + if precision <= 38 { + let value = value.maybe_i128().ok_or_else(|| { + vortex_err!( + "Decimal value {value} cannot fit in Arrow Decimal128 for precision {precision}" + ) + })?; + decimal128_scalar(value, precision, scale) + } else { + decimal256_scalar(value, precision, scale) + } + } + None => { + let data_type = to_data_type_naive(scalar.dtype())?; + Ok(Arc::new(ArrowScalar::new(new_null_array( + &data_type, + SCALAR_ARRAY_LEN, + )))) + } } } +fn decimal128_scalar(value: i128, precision: u8, scale: i8) -> Result, VortexError> { + let array = Decimal128Array::new_scalar(value) + .into_inner() + .with_precision_and_scale(precision, scale)?; + Ok(Arc::new(ArrowScalar::new(array))) +} + +fn decimal256_scalar(value: i256, precision: u8, scale: i8) -> Result, VortexError> { + let array = Decimal256Array::new_scalar(value.into()) + .into_inner() + .with_precision_and_scale(precision, scale)?; + Ok(Arc::new(ArrowScalar::new(array))) +} + /// Convert a [`Utf8Scalar`] to an Arrow [`Datum`]. fn utf8_to_arrow(scalar: Utf8Scalar<'_>) -> Result, VortexError> { value_to_arrow_scalar!(scalar.value(), StringViewArray) @@ -159,7 +191,12 @@ fn map_to_arrow(scalar: MapScalar<'_>) -> Result, VortexError> { let key_array = concat_scalar_arrays(&keys, &key_dtype)?; let value_array = concat_scalar_arrays(&values, &value_dtype)?; - let entries = StructArray::new(fields.clone(), vec![key_array, value_array], None); + let entries = StructArray::try_new_with_length( + fields.clone(), + vec![key_array, value_array], + None, + entries.len(), + )?; let entries_len = entries.len(); let entries_len = i32::try_from(entries_len).map_err(|_| { @@ -271,9 +308,11 @@ mod tests { use std::sync::Arc; use arrow_array::Array; + use arrow_array::Decimal128Array; use arrow_array::Int32Array; use arrow_array::MapArray; use arrow_array::StringViewArray; + use arrow_schema::DataType; use rstest::rstest; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; @@ -435,8 +474,16 @@ mod tests { assert!(result.is_ok()); } + fn assert_arrow_scalar_data_type(scalar: &Scalar, expected: DataType) -> VortexResult<()> { + let datum = scalar.to_arrow_datum()?; + let (array, is_scalar) = datum.get(); + assert!(is_scalar); + assert_eq!(array.data_type(), &expected); + Ok(()) + } + #[test] - fn test_decimal_scalars_to_arrow() { + fn test_decimal_scalars_to_arrow() -> VortexResult<()> { // Test various decimal value types let decimal_dtype = DecimalDType::new(5, 2); @@ -445,37 +492,35 @@ mod tests { decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i8.to_arrow_datum().is_ok()); + assert_arrow_scalar_data_type(&scalar_i8, DataType::Decimal128(5, 2))?; let scalar_i16 = Scalar::decimal( DecimalValue::I16(10000), decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i16.to_arrow_datum().is_ok()); + assert_arrow_scalar_data_type(&scalar_i16, DataType::Decimal128(5, 2))?; let scalar_i32 = Scalar::decimal( DecimalValue::I32(99999), decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i32.to_arrow_datum().is_ok()); + assert_arrow_scalar_data_type(&scalar_i32, DataType::Decimal128(5, 2))?; let scalar_i64 = Scalar::decimal( DecimalValue::I64(99999), decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i64.to_arrow_datum().is_ok()); + assert_arrow_scalar_data_type(&scalar_i64, DataType::Decimal128(5, 2))?; let scalar_i128 = Scalar::decimal( DecimalValue::I128(99999), decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i128.to_arrow_datum().is_ok()); - - // Test i256 + assert_arrow_scalar_data_type(&scalar_i128, DataType::Decimal128(5, 2))?; let value_i256 = i256::from_i128(99999); let scalar_i256 = Scalar::decimal( @@ -483,15 +528,64 @@ mod tests { decimal_dtype, Nullability::NonNullable, ); - assert!(scalar_i256.to_arrow_datum().is_ok()); + assert_arrow_scalar_data_type(&scalar_i256, DataType::Decimal128(5, 2))?; + + Ok(()) + } + + #[test] + fn decimal_i64_with_wide_precision_exports_decimal256() -> VortexResult<()> { + let scalar = Scalar::decimal( + DecimalValue::I64(1), + DecimalDType::new(39, 0), + Nullability::NonNullable, + ); + + assert_arrow_scalar_data_type(&scalar, DataType::Decimal256(39, 0)) + } + + #[test] + fn decimal_i256_with_narrow_precision_exports_decimal128() -> VortexResult<()> { + let scalar = Scalar::decimal( + DecimalValue::I256(i256::from_i128(1234)), + DecimalDType::new(4, 2), + Nullability::NonNullable, + ); + + assert_arrow_scalar_data_type(&scalar, DataType::Decimal128(4, 2)) } #[test] - fn test_null_decimal_to_arrow() { + fn test_null_decimal_to_arrow() -> VortexResult<()> { let decimal_dtype = DecimalDType::new(10, 2); let scalar = Scalar::null(DType::Decimal(decimal_dtype, Nullability::Nullable)); - let result = scalar.to_arrow_datum(); - assert!(result.is_ok()); + assert_arrow_scalar_data_type(&scalar, DataType::Decimal128(10, 2))?; + + let decimal_dtype = DecimalDType::new(39, 2); + let scalar = Scalar::null(DType::Decimal(decimal_dtype, Nullability::Nullable)); + assert_arrow_scalar_data_type(&scalar, DataType::Decimal256(39, 2)) + } + + #[test] + fn decimal_scalar_to_arrow_preserves_precision_and_scale() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(12, 3); + let scalar = Scalar::decimal( + DecimalValue::I128(12345), + decimal_dtype, + Nullability::NonNullable, + ); + + let datum = scalar.to_arrow_datum()?; + let (array, is_scalar) = datum.get(); + assert!(is_scalar); + let decimal = array + .as_any() + .downcast_ref::() + .expect("decimal scalar should convert to Decimal128"); + assert_eq!(decimal.precision(), 12); + assert_eq!(decimal.scale(), 3); + + Ok(()) } #[test] @@ -546,7 +640,102 @@ mod tests { } #[test] - #[should_panic(expected = "struct scalar conversion")] + fn map_decimal_scalar_to_arrow_preserves_decimal_type() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(9, 2); + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Decimal(decimal_dtype, Nullability::Nullable), + false, + Nullability::NonNullable, + )?; + let scalar = Scalar::try_map( + dtype, + [( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::decimal( + DecimalValue::I256(i256::from_i128(12345)), + decimal_dtype, + Nullability::Nullable, + ), + )], + )?; + + let datum = scalar.to_arrow_datum()?; + let (array, is_scalar) = datum.get(); + assert!(is_scalar); + let map = array + .as_any() + .downcast_ref::() + .expect("map scalar should convert to MapArray"); + assert_eq!(map.values().data_type(), &DataType::Decimal128(9, 2)); + + let wide_decimal_dtype = DecimalDType::new(39, 2); + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + DType::Decimal(wide_decimal_dtype, Nullability::Nullable), + false, + Nullability::NonNullable, + )?; + let scalar = Scalar::try_map( + dtype, + [( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::decimal( + DecimalValue::I64(12345), + wide_decimal_dtype, + Nullability::Nullable, + ), + )], + )?; + + let datum = scalar.to_arrow_datum()?; + let (array, is_scalar) = datum.get(); + assert!(is_scalar); + let map = array + .as_any() + .downcast_ref::() + .expect("map scalar should convert to MapArray"); + assert_eq!(map.values().data_type(), &DataType::Decimal256(39, 2)); + + Ok(()) + } + + #[test] + fn map_scalar_with_unsupported_nested_value_errors_without_panic() -> VortexResult<()> { + let struct_dtype = DType::Struct( + StructFields::from_iter([( + "field1", + FieldDType::from(DType::Primitive(PType::I32, Nullability::NonNullable)), + )]), + Nullability::NonNullable, + ); + let dtype = DType::map( + DType::Primitive(PType::I32, Nullability::NonNullable), + struct_dtype.clone(), + false, + Nullability::NonNullable, + )?; + let scalar = Scalar::try_map( + dtype, + [( + Scalar::primitive(1i32, Nullability::NonNullable), + Scalar::struct_( + struct_dtype, + vec![Scalar::primitive(42i32, Nullability::NonNullable)], + ), + )], + )?; + + let error = scalar + .to_arrow_datum() + .err() + .expect("unsupported nested map value should error"); + assert!(error.to_string().contains("struct scalar conversion")); + + Ok(()) + } + + #[test] fn test_struct_scalar_to_arrow_todo() { let struct_dtype = DType::Struct( StructFields::from_iter([( @@ -560,11 +749,14 @@ mod tests { struct_dtype, vec![Scalar::primitive(42i32, Nullability::NonNullable)], ); - struct_scalar.to_arrow_datum().unwrap(); + let error = struct_scalar + .to_arrow_datum() + .err() + .expect("struct scalar should error"); + assert!(error.to_string().contains("struct scalar conversion")); } #[test] - #[should_panic(expected = "list scalar conversion")] fn test_list_scalar_to_arrow_todo() { let element_dtype = Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)); let list_scalar = Scalar::list( @@ -576,11 +768,14 @@ mod tests { Nullability::NonNullable, ); - list_scalar.to_arrow_datum().unwrap(); + let error = list_scalar + .to_arrow_datum() + .err() + .expect("list scalar should error"); + assert!(error.to_string().contains("list scalar conversion")); } #[test] - #[should_panic(expected = "Cannot convert extension scalar")] fn test_non_temporal_extension_to_arrow_todo() { #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] struct SomeExt; @@ -618,7 +813,15 @@ mod tests { Scalar::primitive(42i32, Nullability::NonNullable), ); - scalar.to_arrow_datum().unwrap(); + let error = scalar + .to_arrow_datum() + .err() + .expect("non-temporal extension scalar should error"); + assert!( + error + .to_string() + .contains("Cannot convert extension scalar") + ); } #[rstest] diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index 87a08a548ea..7f19c3af3c3 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -63,6 +63,7 @@ use vortex_session::registry::Id; use crate::FromArrowArray; use crate::IntoVortexArray; +use crate::convert::map_from_arrow_parts; use crate::convert::nulls; use crate::convert::remove_nulls; use crate::dtype::TryFromArrowType; @@ -620,8 +621,17 @@ impl ArrowSession { let validity = nulls(list.nulls(), field.is_nullable())?; Ok(ListViewArray::try_new(elements, offsets, sizes, validity)?.into_array()) } - DataType::Map(..) => { - vortex_bail!("Arrow MapArray conversion is not yet supported") + DataType::Map(entries_field, keys_sorted) => { + let map = array.as_map(); + let entries_array: ArrowArrayRef = Arc::new(map.entries().clone()); + let entries = self.from_arrow_array(entries_array, entries_field.as_ref())?; + map_from_arrow_parts( + entries, + map.offsets(), + map.nulls(), + *keys_sorted, + field.is_nullable(), + ) } _ => ArrayRef::from_arrow(array.as_ref(), field.is_nullable()), } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 06af9e25e8c..b11c84e55aa 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -30,8 +30,10 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::assert_arrays_eq; +use vortex_array::builders::MapBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::MapDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; @@ -1825,15 +1827,29 @@ async fn test_writer_with_complex_types() -> VortexResult<()> { /// Write `array` with list decomposition forced on (through the full compress/zone pipeline) and /// read the whole thing back. async fn write_read_roundtrip(array: ArrayRef) -> VortexResult { + write_read_roundtrip_with_layout(array, true).await +} + +async fn write_read_roundtrip_with_layout( + array: ArrayRef, + use_list_layout: bool, +) -> VortexResult { let strategy = crate::strategy::WriteStrategyBuilder::default() .with_list_layout() .build(); let mut buf = ByteBufferMut::empty(); - SESSION - .write_options() - .with_strategy(strategy) - .write(&mut buf, array.to_array_stream()) - .await?; + if use_list_layout { + SESSION + .write_options() + .with_strategy(strategy) + .write(&mut buf, array.to_array_stream()) + .await?; + } else { + SESSION + .write_options() + .write(&mut buf, array.to_array_stream()) + .await?; + } SESSION .open_options() .open_buffer(buf)? @@ -1867,6 +1883,71 @@ async fn nested_list_of_list_roundtrip() -> VortexResult<()> { Ok(()) } +type MapEntryFixture<'a> = (i32, Option<&'a str>); +type MapRowFixture<'a> = Option>>; + +fn map_array_from_rows(rows: &[MapRowFixture<'_>], keys_sorted: bool) -> VortexResult { + let map_dtype = MapDType::try_new( + DType::Primitive(I32, Nullability::NonNullable), + DType::Utf8(Nullability::Nullable), + keys_sorted, + )?; + let dtype = DType::Map(map_dtype.clone(), Nullability::Nullable); + let mut builder = + MapBuilder::::with_capacity(map_dtype, Nullability::Nullable, rows.len()); + + for row in rows { + let scalar = match row { + Some(entries) => { + let entries = entries + .iter() + .map(|(key, value)| { + let key = Scalar::primitive(*key, Nullability::NonNullable); + let value = value.map_or_else( + || Scalar::null(DType::Utf8(Nullability::Nullable)), + |value| Scalar::utf8(value, Nullability::Nullable), + ); + (key, value) + }) + .collect::>(); + Scalar::try_map(dtype.clone(), entries)? + } + None => Scalar::null(dtype.clone()), + }; + builder.append_value(scalar.as_map())?; + } + + Ok(builder.finish_into_map().into_array()) +} + +/// A struct containing a Map column crosses both the default flat writer and the list layout +/// strategy without changing map nullability, empty rows, duplicate keys, or scalar values. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn struct_with_map_column_roundtrip() -> VortexResult<()> { + for use_list_layout in [false, true] { + let maps = map_array_from_rows( + &[ + Some(vec![(1, Some("one")), (2, None)]), + None, + Some(vec![]), + Some(vec![(1, Some("dup-old")), (1, Some("dup-new"))]), + ], + false, + )?; + let st = StructArray::from_fields(&[ + ("id", buffer![10i32, 20, 30, 40].into_array()), + ("attrs", maps), + ])? + .into_array(); + + let result = write_read_roundtrip_with_layout(st.clone(), use_list_layout).await?; + assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx()); + } + + Ok(()) +} + /// A `struct<{ items: list>? }>` column round-trips, exercising list decomposition /// recursing into struct decomposition (list `elements` are structs) plus a nullable list validity /// child.