Skip to content

Execute until arrays match fast path or canonical before exporting to Arrow #8935

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mhk197 merged 1 commit into from
Jul 27, 2026
Merged

Execute until arrays match fast path or canonical before exporting to Arrow #8935

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 53 additions & 1 deletion vortex-arrow/src/executor/byte.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ use vortex_array::ArrayRef;
use vortex_array::ArrayView;
use vortex_array::Canonical;
use vortex_array::ExecutionCtx;
use vortex_array::arrays::Chunked;
use vortex_array::arrays::Constant;
use vortex_array::arrays::VarBin;
use vortex_array::arrays::VarBinViewArray;
use vortex_array::arrays::varbin::VarBinArraySlotsExt;
Expand All @@ -22,13 +24,29 @@ use vortex_array::dtype::DType;
use vortex_array::dtype::NativePType;
use vortex_array::dtype::Nullability;
use vortex_array::dtype::PType;
use vortex_array::matcher::Matcher;
use vortex_error::VortexError;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;

use crate::byte_view::execute_varbinview_to_arrow;
use crate::executor::validity::to_arrow_null_buffer;

/// Matches the encodings [`to_arrow_byte_array`] requires for export.
///
/// `Chunked` and `Constant` are matched to stop execution before it destroys them: they have
/// specialized `append_to_builder` impls (chunk-wise append, scalar repeat) that the builder
/// fallback exploits.
struct ArrowByteExportable;

impl Matcher for ArrowByteExportable {
type Match<'a> = &'a ArrayRef;

fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
(array.is::<VarBin>() || array.is::<Chunked>() || array.is::<Constant>()).then_some(array)
}
}

/// Convert a Vortex array into an Arrow GenericBinaryArray.
pub(super) fn to_arrow_byte_array<T: ByteArrayType>(
array: ArrayRef,
Expand Down Expand Up @@ -57,7 +75,9 @@ where
return arrow_cast::cast(&binary_view, &T::DATA_TYPE).map_err(VortexError::from);
}

// If the Vortex array is already in VarBin format, we can directly convert it.
let array = array.execute_until::<ArrowByteExportable>(ctx)?;

// If the Vortex array is in VarBin format, we can directly convert it.
if let Some(array) = array.as_opt::<VarBin>() {
return varbin_to_byte_array::<T>(array, ctx);
}
Expand Down Expand Up @@ -106,15 +126,47 @@ mod tests {
use vortex_array::IntoArray;
use vortex_array::VortexSessionExecute;
use vortex_array::array_session;
use vortex_array::arrays::BoolArray;
use vortex_array::arrays::PrimitiveArray;
use vortex_array::arrays::VarBinArray;
use vortex_array::arrays::VarBinViewArray;
use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt;
use vortex_array::dtype::DType;
use vortex_array::dtype::Nullability;
use vortex_array::scalar_fn::EmptyOptions;
use vortex_array::scalar_fn::fns::mask::Mask as MaskFn;
use vortex_error::VortexResult;
use vortex_mask::Mask;

use crate::ArrowSessionExt;

#[test]
fn mask_wrapped_varbin_exports() -> VortexResult<()> {
let session = array_session();
let mut ctx = session.create_execution_ctx();

let varbin = VarBinArray::from_vec(
vec!["hello", "world", "vortex"],
DType::Utf8(Nullability::NonNullable),
);
let mask = BoolArray::from_iter([true, false, true]);
let masked =
MaskFn.try_new_array(3, EmptyOptions, [varbin.into_array(), mask.into_array()])?;

let field = Field::new("s", DataType::Utf8, true);
let arrow = session
.arrow()
.execute_arrow(masked, Some(&field), &mut ctx)?;

let strings = arrow.as_string::<i32>();
assert_eq!(strings.len(), 3);
assert!(!strings.is_null(0));
assert!(strings.is_null(1));
assert_eq!(strings.value(0), "hello");
assert_eq!(strings.value(2), "vortex");
Ok(())
}

fn make_utf8_array() -> VarBinViewArray {
VarBinViewArray::from_iter_str(["hello", "world", "this is a longer string for testing"])
}
Expand Down
38 changes: 38 additions & 0 deletions vortex-arrow/src/executor/dictionary.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,18 +18,32 @@ use vortex_array::arrays::ConstantArray;
use vortex_array::arrays::Dict;
use vortex_array::arrays::DictArray;
use vortex_array::arrays::dict::DictArraySlotsExt;
use vortex_array::matcher::Matcher;
use vortex_error::VortexError;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;

use crate::ArrowArrayExecutor;

/// Matches the encodings [`to_arrow_dictionary`] requires for export.
struct ArrowDictExportable;

impl Matcher for ArrowDictExportable {
type Match<'a> = &'a ArrayRef;

fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
(array.is::<Dict>() || array.is::<Constant>()).then_some(array)
}
}

pub(super) fn to_arrow_dictionary(
array: ArrayRef,
codes_type: &DataType,
values_type: &DataType,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrowArrayRef> {
let array = array.execute_until::<ArrowDictExportable>(ctx)?;

let array = match array.try_downcast::<Dict>() {
Ok(dict) => return dict_to_dict(dict, codes_type, values_type, ctx),
Err(array) => array,
Expand Down Expand Up @@ -141,18 +155,23 @@ mod tests {
use std::sync::Arc;

use arrow_array::DictionaryArray as ArrowDictArray;
use arrow_array::StringArray;
use arrow_array::types::UInt8Type;
use arrow_array::types::UInt32Type;
use arrow_schema::DataType;
use rstest::rstest;
use vortex_array::IntoArray;
use vortex_array::VortexSessionExecute;
use vortex_array::array_session;
use vortex_array::arrays::BoolArray;
use vortex_array::arrays::PrimitiveArray;
use vortex_array::arrays::VarBinViewArray;
use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt;
use vortex_array::dtype::DType;
use vortex_array::dtype::Nullability::Nullable;
use vortex_array::scalar::Scalar;
use vortex_array::scalar_fn::EmptyOptions;
use vortex_array::scalar_fn::fns::mask::Mask;
use vortex_buffer::buffer;
use vortex_error::VortexResult;

Expand Down Expand Up @@ -224,4 +243,23 @@ mod tests {
assert_eq!(expected.as_ref(), actual.as_ref());
Ok(())
}

#[test]
fn mask_wrapped_dict_exports() -> VortexResult<()> {
// Dictionary behind a lazy `mask` scalar-fn — the shape a scan produces when a row
// mask is applied to a dict-encoded column.
let dict = DictArray::try_new(
buffer![0u8, 1, 0].into_array(),
VarBinViewArray::from_iter_str(["a", "b"]).into_array(),
)?;
let mask = BoolArray::from_iter([true, false, true]);
let masked = Mask.try_new_array(3, EmptyOptions, [dict.into_array(), mask.into_array()])?;

let actual = execute(masked, &dict_type(DataType::UInt8, DataType::Utf8))?;
let flat = arrow_cast::cast(&actual, &DataType::Utf8)?;

let expected = StringArray::from(vec![Some("a"), None, Some("a")]);
assert_eq!(flat.as_ref(), &expected as &dyn arrow_array::Array);
Ok(())
}
}
131 changes: 114 additions & 17 deletions vortex-arrow/src/executor/list.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ use vortex_array::builtins::ArrayBuiltins;
use vortex_array::dtype::DType;
use vortex_array::dtype::NativePType;
use vortex_array::dtype::Nullability;
use vortex_array::matcher::Matcher;
use vortex_buffer::BufferMut;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
Expand All @@ -33,16 +34,29 @@ use vortex_error::vortex_ensure;
use crate::executor::validity::to_arrow_null_buffer;
use crate::session::ArrowSessionExt;

/// Matches the encodings [`to_arrow_list`] requires for export.
struct ArrowListExportable;

impl Matcher for ArrowListExportable {
type Match<'a> = &'a ArrayRef;

fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
(array.is::<List>() || array.is::<Chunked>() || array.is::<ListView>()).then_some(array)
}
}

#[allow(rustdoc::broken_intra_doc_links)]
/// Convert a Vortex VarBinArray into an Arrow [`GenericListArray`](arrow_array:array::GenericListArray).
pub(super) fn to_arrow_list<O: OffsetSizeTrait + NativePType>(
array: ArrayRef,
elements_field: &FieldRef,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrowArrayRef> {
let array = array.execute_until::<ArrowListExportable>(ctx)?;

// If the Vortex array is already in List format, we can directly convert it.
if let Some(array) = array.as_opt::<List>() {
return list_to_list::<O>(&array.into_owned(), elements_field, ctx);
if let Some(list) = array.as_opt::<List>() {
return list_to_list::<O>(&list.into_owned(), elements_field, ctx);
}

// Converting each chunk individually, then using the fast concat logic from arrow
Expand All @@ -56,24 +70,15 @@ pub(super) fn to_arrow_list<O: OffsetSizeTrait + NativePType>(
return Ok(arrow_select::concat::concat(&refs)?);
}

// If the Vortex array is a ListViewArray, rebuild to ZCTL if needed and convert.
let array = match array.try_downcast::<ListView>() {
Ok(array) => {
let zctl = if array.is_zero_copy_to_list() {
array
} else {
array.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)?
};
return list_view_zctl::<O>(zctl, elements_field, ctx);
}
Err(a) => a,
};

// Otherwise, we execute the array to become a ListViewArray, then rebuild to ZCTL.
// Otherwise the array is canonical: a ListViewArray, which we rebuild to ZCTL if needed.
// Note: arrow_cast::cast supports ListView → List (apache/arrow-rs#8735), but it
// unconditionally uses take. Our rebuild uses a heuristic that picks list-by-list
// for large lists, which avoids materializing a large index buffer.
let list_view = array.execute::<ListViewArray>(ctx)?;
let list_view = array
.as_opt::<ListView>()
.vortex_expect("Must be ListView from Matcher")
.into_owned();

let zctl = if list_view.is_zero_copy_to_list() {
list_view
} else {
Expand Down Expand Up @@ -213,15 +218,21 @@ mod tests {
use vortex_array::Canonical;
use vortex_array::IntoArray;
use vortex_array::VortexSessionExecute;
use vortex_array::arrays::BoolArray;
use vortex_array::arrays::PrimitiveArray;
use vortex_array::arrays::SliceArray;
use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt;
use vortex_array::dtype::DType;
use vortex_array::dtype::Nullability::NonNullable;
use vortex_array::scalar_fn::EmptyOptions;
use vortex_array::scalar_fn::fns::mask::Mask;
use vortex_array::validity::Validity;
use vortex_buffer::buffer;
use vortex_error::VortexResult;
use vortex_session::VortexSession;

use crate::ArrowArrayExecutor;
use crate::executor::list::ListArray;
use crate::executor::list::ListViewArray;

/// A shared session for these list-executor tests, used to create execution contexts.
Expand Down Expand Up @@ -364,6 +375,92 @@ mod tests {
Ok(())
}

#[test]
fn slice_wrapped_list_exports() -> VortexResult<()> {
let mut ctx = SESSION.create_execution_ctx();
// Lists [[1, 2], [3], [4, 5, 6]] behind a lazy Slice wrapper selecting rows 1..3.
let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable);
let elements_ptr = elements.as_slice::<i32>().as_ptr();
let offsets = PrimitiveArray::new(buffer![0i32, 2, 3, 6], Validity::NonNullable);
let list = ListArray::try_new(
elements.into_array(),
offsets.into_array(),
Validity::NonNullable,
)?;
let sliced = SliceArray::new(list.into_array(), 1..3).into_array();

let field = Field::new("item", DataType::Int32, false);
let arrow_dt = DataType::List(field.into());
let arrow_array = sliced.execute_arrow(Some(&arrow_dt), &mut ctx)?;

let arrow_list = arrow_array
.as_any()
.downcast_ref::<GenericListArray<i32>>()
.unwrap();
assert_eq!(arrow_list.len(), 2);
let first = arrow_list.value(0);
let first_vals = first.as_any().downcast_ref::<Int32Array>().unwrap();
assert_eq!(first_vals.values(), &[3]);
let second = arrow_list.value(1);
let second_vals = second.as_any().downcast_ref::<Int32Array>().unwrap();
assert_eq!(second_vals.values(), &[4, 5, 6]);

// The conversion shares the elements buffer regardless of which path resolves the
// Slice (revealed List or zero-copy-to-list ListView).
let values = arrow_list
.values()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
assert_eq!(values.values().as_ptr(), elements_ptr);
Ok(())
}

#[test]
fn mask_wrapped_list_exports() -> VortexResult<()> {
let mut ctx = SESSION.create_execution_ctx();
// Lists [[1, 2], [3], [4, 5, 6]] behind a lazy `mask` scalar-fn nulling out row 1 —
// the shape a scan produces when a row mask is applied to a List-encoded column.
let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable);
let elements_ptr = elements.as_slice::<i32>().as_ptr();
let offsets = PrimitiveArray::new(buffer![0i32, 2, 3, 6], Validity::NonNullable);
let list = ListArray::try_new(
elements.into_array(),
offsets.into_array(),
Validity::NonNullable,
)?;
let mask = BoolArray::from_iter([true, false, true]);
let masked = Mask.try_new_array(3, EmptyOptions, [list.into_array(), mask.into_array()])?;

let field = Field::new("item", DataType::Int32, false);
let arrow_dt = DataType::List(field.into());
let arrow_array = masked.execute_arrow(Some(&arrow_dt), &mut ctx)?;

let arrow_list = arrow_array
.as_any()
.downcast_ref::<GenericListArray<i32>>()
.unwrap();
assert_eq!(arrow_list.len(), 3);
assert!(!arrow_list.is_null(0));
assert!(arrow_list.is_null(1));
assert!(!arrow_list.is_null(2));
let first = arrow_list.value(0);
let first_vals = first.as_any().downcast_ref::<Int32Array>().unwrap();
assert_eq!(first_vals.values(), &[1, 2]);
let third = arrow_list.value(2);
let third_vals = third.as_any().downcast_ref::<Int32Array>().unwrap();
assert_eq!(third_vals.values(), &[4, 5, 6]);

// Masking only touches validity, so the conversion still shares the elements buffer.
let values = arrow_list
.values()
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
assert_eq!(values.values().as_ptr(), elements_ptr);
Ok(())
}

#[test]
fn test_to_arrow_list_empty_zctl() -> VortexResult<()> {
let mut ctx = SESSION.create_execution_ctx();
Expand Down
Loading
Loading