Skip to content

Commit 28b0e32

Browse files
committed
Add Arrow Map roundtrip support
Signed-off-by: Adam Gutglick <adam@spiraldb.com>
1 parent 2ce8917 commit 28b0e32

12 files changed

Lines changed: 1072 additions & 49 deletions

File tree

‎vortex-arrow/src/convert.rs‎

Lines changed: 79 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ use arrow_array::GenericByteArray;
1313
use arrow_array::GenericByteViewArray;
1414
use arrow_array::GenericListArray;
1515
use arrow_array::GenericListViewArray;
16+
use arrow_array::MapArray as ArrowMapArray;
1617
use arrow_array::NullArray as ArrowNullArray;
1718
use arrow_array::OffsetSizeTrait;
1819
use arrow_array::PrimitiveArray as ArrowPrimitiveArray;
@@ -65,16 +66,22 @@ use vortex_array::arrays::DictArray;
6566
use vortex_array::arrays::FixedSizeListArray;
6667
use vortex_array::arrays::ListArray;
6768
use vortex_array::arrays::ListViewArray;
69+
use vortex_array::arrays::MapArray;
6870
use vortex_array::arrays::NullArray;
6971
use vortex_array::arrays::PrimitiveArray;
72+
use vortex_array::arrays::Struct;
7073
use vortex_array::arrays::StructArray;
7174
use vortex_array::arrays::TemporalArray;
7275
use vortex_array::arrays::VarBinArray;
7376
use vortex_array::arrays::VarBinViewArray;
77+
use vortex_array::arrays::struct_::StructArrayExt;
7478
use vortex_array::dtype::DType;
7579
use vortex_array::dtype::DecimalDType;
80+
use vortex_array::dtype::FieldNames;
7681
use vortex_array::dtype::IntegerPType;
82+
use vortex_array::dtype::MapDType;
7783
use vortex_array::dtype::NativePType;
84+
use vortex_array::dtype::Nullability;
7885
use vortex_array::dtype::PType;
7986
use vortex_array::dtype::i256;
8087
use vortex_array::extension::datetime::TimeUnit;
@@ -477,6 +484,77 @@ impl FromArrowArray<&ArrowFixedSizeListArray> for ArrayRef {
477484
}
478485
}
479486

487+
pub(crate) fn map_from_arrow_parts(
488+
entries: ArrayRef,
489+
offsets: &OffsetBuffer<i32>,
490+
arrow_nulls: Option<&NullBuffer>,
491+
keys_sorted: bool,
492+
nullable: bool,
493+
) -> VortexResult<ArrayRef> {
494+
let DType::Struct(struct_dtype, Nullability::NonNullable) = entries.dtype() else {
495+
vortex_bail!(
496+
"Arrow map entries must import as non-nullable struct, got {}",
497+
entries.dtype()
498+
);
499+
};
500+
vortex_ensure!(
501+
struct_dtype.nfields() == 2,
502+
"Arrow map entries struct must contain exactly two fields"
503+
);
504+
505+
let key_dtype = struct_dtype
506+
.field_by_index(0)
507+
.ok_or_else(|| vortex_err!("Arrow map entries struct missing key field"))?;
508+
vortex_ensure!(
509+
!key_dtype.is_nullable(),
510+
"Arrow map key field must be non-nullable"
511+
);
512+
let value_dtype = struct_dtype
513+
.field_by_index(1)
514+
.ok_or_else(|| vortex_err!("Arrow map entries struct missing value field"))?;
515+
let map_dtype = MapDType::try_new(key_dtype, value_dtype, keys_sorted)?;
516+
let entries_struct = entries.as_::<Struct>();
517+
let entries = StructArray::try_new(
518+
FieldNames::from(["key", "value"]),
519+
vec![
520+
entries_struct.unmasked_field(0).clone(),
521+
entries_struct.unmasked_field(1).clone(),
522+
],
523+
entries.len(),
524+
Validity::NonNullable,
525+
)?
526+
.into_array();
527+
528+
let len = offsets.len() - 1;
529+
let sizes = Buffer::<i32>::from_iter(offsets.windows(2).map(|window| window[1] - window[0]))
530+
.into_array();
531+
let offsets = offsets.inner().slice(0, len).into_array();
532+
let validity = nulls(arrow_nulls, nullable)?;
533+
let entries = ListViewArray::try_new(entries, offsets, sizes, validity)?;
534+
// SAFETY: Arrow Map offsets are sorted list offsets. Vortex sizes are adjacent offset
535+
// differences, so every view's end equals the next view's start. Sliced Arrow arrays may
536+
// leave leading or trailing entries unreferenced, which zero-copy-to-list permits.
537+
let entries = unsafe { entries.with_zero_copy_to_list(true) };
538+
539+
Ok(MapArray::try_new(map_dtype, entries)?.into_array())
540+
}
541+
542+
impl FromArrowArray<&ArrowMapArray> for ArrayRef {
543+
fn from_arrow(array: &ArrowMapArray, nullable: bool) -> VortexResult<Self> {
544+
let DataType::Map(_, keys_sorted) = array.data_type() else {
545+
vortex_panic!("Invalid data type for MapArray: {}", array.data_type());
546+
};
547+
let entries = Self::from_arrow(array.entries(), false)?;
548+
map_from_arrow_parts(
549+
entries,
550+
array.offsets(),
551+
array.nulls(),
552+
*keys_sorted,
553+
nullable,
554+
)
555+
}
556+
}
557+
480558
impl FromArrowArray<&ArrowNullArray> for ArrayRef {
481559
fn from_arrow(value: &ArrowNullArray, nullable: bool) -> VortexResult<Self> {
482560
vortex_ensure!(
@@ -546,9 +624,7 @@ impl FromArrowArray<&dyn ArrowArray> for ArrayRef {
546624
DataType::ListView(_) => Self::from_arrow(array.as_list_view::<i32>(), nullable),
547625
DataType::LargeListView(_) => Self::from_arrow(array.as_list_view::<i64>(), nullable),
548626
DataType::FixedSizeList(..) => Self::from_arrow(array.as_fixed_size_list(), nullable),
549-
DataType::Map(..) => {
550-
vortex_bail!("Arrow MapArray conversion is not yet supported")
551-
}
627+
DataType::Map(..) => Self::from_arrow(array.as_map(), nullable),
552628
DataType::Null => Self::from_arrow(as_null_array(array), nullable),
553629
DataType::Timestamp(u, _) => match u {
554630
ArrowTimeUnit::Second => {

0 commit comments

Comments
 (0)