From c989232ed61185ec7fbc81c742d7f49b13aa0014 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 23 Oct 2026 14:16:04 +0100 Subject: [PATCH] Dispatch filter take of fixed-width values on byte width The filter take kernel copies values and does not inspect them, but it was compiled for each of the 17 primitive and decimal value types. The primitive and decimal paths now reinterpret the values buffer by byte width, so six kernels (1, 2, 4, 8, 16 and 32 bytes) serve all types. The reinterpretation does not copy data. vortex-array release __text: 14,012,988 -> 13,808,764 bytes (-204 KB). Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5.5 --- .../arrays/filter/execute/take/fixed_width.rs | 118 ++++++++++++------ .../src/arrays/filter/execute/take/tests.rs | 58 +++++++++ 2 files changed, 140 insertions(+), 36 deletions(-) diff --git a/vortex-array/src/arrays/filter/execute/take/fixed_width.rs b/vortex-array/src/arrays/filter/execute/take/fixed_width.rs index 7a66aa8320b..1f3fe8209db 100644 --- a/vortex-array/src/arrays/filter/execute/take/fixed_width.rs +++ b/vortex-array/src/arrays/filter/execute/take/fixed_width.rs @@ -5,7 +5,9 @@ use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; +use vortex_error::vortex_panic; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -22,11 +24,11 @@ use crate::arrays::Filter; use crate::arrays::PrimitiveArray; use crate::arrays::decimal::DecimalArrayExt; use crate::arrays::filter::FilterArraySlotsExt; +use crate::buffer::BufferHandle; use crate::dtype::IntegerPType; +use crate::dtype::i256; use crate::executor::ExecutionCtx; -use crate::match_each_decimal_value_type; use crate::match_each_integer_ptype; -use crate::match_each_native_ptype; use crate::validity::Validity; #[inline] @@ -36,21 +38,17 @@ pub(in crate::arrays::filter) fn take_primitive( ctx: &mut ExecutionCtx, ) -> VortexResult { let child = array.child().clone().execute::(ctx)?; - let child_validity = child.validity()?; - match_each_native_ptype!(child.ptype(), |T| { - let child_buf = child.to_buffer::(); - match_each_integer_ptype!(indices.ptype(), |P| { - let (taken, output_validity) = take_fixed_width::( - child_buf, - child_validity, - array.filter_mask(), - indices, - ctx, - )?; - // SAFETY: Take operation validated all the parts - Ok(unsafe { PrimitiveArray::new_unchecked(taken, output_validity) }.into_array()) - }) - }) + let ptype = child.ptype(); + let (taken, output_validity) = take_fixed_width_bytes( + child.buffer_handle().to_host_sync(), + ptype.byte_width(), + child.validity()?, + array.filter_mask(), + indices, + ctx, + )?; + + Ok(PrimitiveArray::from_byte_buffer(taken, ptype, output_validity).into_array()) } #[inline] @@ -60,28 +58,76 @@ pub(in crate::arrays::filter) fn take_decimal( ctx: &mut ExecutionCtx, ) -> VortexResult { let child = array.child().clone().execute::(ctx)?; - let decimal_dtype = child.decimal_dtype(); - let child_validity = child.validity()?; - match_each_decimal_value_type!(child.values_type(), |T| { - let child_buf = child.buffer::(); - match_each_integer_ptype!(indices.ptype(), |P| { - let (taken, output_validity) = take_fixed_width::( - child_buf, - child_validity, - array.filter_mask(), - indices, - ctx, - )?; - // SAFETY: Valid ranks copy existing decimal values, null ranks write default placeholders that - // are hidden by output validity, and the output validity was built for the take length. - Ok( - unsafe { DecimalArray::new_unchecked(taken, decimal_dtype, output_validity) } - .into_array(), - ) - }) + let values_type = child.values_type(); + let (taken, output_validity) = take_fixed_width_bytes( + child.buffer_handle().to_host_sync(), + values_type.byte_width(), + child.validity()?, + array.filter_mask(), + indices, + ctx, + )?; + + // SAFETY: Valid ranks copy existing decimal values, null ranks write default placeholders that + // are hidden by output validity, and the output validity was built for the take length. + Ok(unsafe { + DecimalArray::new_unchecked_handle( + BufferHandle::new_host(taken), + values_type, + child.decimal_dtype(), + output_validity, + ) + } + .into_array()) +} + +/// Takes values of `byte_width` bytes each. +/// +/// The take copies values and does not inspect them, so all value types of one width share one +/// kernel. +fn take_fixed_width_bytes( + child: ByteBuffer, + byte_width: usize, + child_validity: Validity, + filter: &Mask, + indices: &PrimitiveArray, + ctx: &mut ExecutionCtx, +) -> VortexResult<(ByteBuffer, Validity)> { + match_each_integer_ptype!(indices.ptype(), |P| { + match byte_width { + 1 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), + 2 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), + 4 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), + 8 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), + 16 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), + 32 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), + _ => vortex_panic!("Unsupported fixed width of {byte_width} bytes"), + } }) } +fn take_fixed_width_as( + child: ByteBuffer, + child_validity: Validity, + filter: &Mask, + indices: &PrimitiveArray, + ctx: &mut ExecutionCtx, +) -> VortexResult<(ByteBuffer, Validity)> +where + W: Copy + Default, + P: IntegerPType, +{ + let (taken, output_validity) = take_fixed_width::( + Buffer::from_byte_buffer(child), + child_validity, + filter, + indices, + ctx, + )?; + + Ok((taken.into_byte_buffer(), output_validity)) +} + fn take_fixed_width( child: Buffer, child_validity: Validity, diff --git a/vortex-array/src/arrays/filter/execute/take/tests.rs b/vortex-array/src/arrays/filter/execute/take/tests.rs index a1e75d7652f..d869cde3f09 100644 --- a/vortex-array/src/arrays/filter/execute/take/tests.rs +++ b/vortex-array/src/arrays/filter/execute/take/tests.rs @@ -3,6 +3,7 @@ use std::sync::LazyLock; +use rstest::rstest; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -26,6 +27,8 @@ use crate::arrays::dict::TakeExecuteAdaptor; use crate::assert_arrays_eq; use crate::dtype::DecimalDType; use crate::dtype::FieldNames; +use crate::dtype::half::f16; +use crate::dtype::i256; use crate::executor::ExecutionCtx; use crate::kernel::ExecuteParentKernel; use crate::validity::Validity; @@ -420,6 +423,61 @@ fn test_take_execute_kernel_handles_nullable_decimal_filter_child() -> VortexRes Ok(()) } +/// Each case has values of a different width, which the take copies as raw bytes. +#[rstest] +#[case::f16( + PrimitiveArray::from_iter([1.0f32, 2.0, 3.0, 4.0, 5.0].map(f16::from_f32)).into_array(), + PrimitiveArray::from_option_iter([Some(f16::from_f32(4.0)), None, Some(f16::from_f32(1.0))]) + .into_array(), +)] +#[case::f32( + PrimitiveArray::from_iter([1.5f32, 2.5, 3.5, 4.5, 5.5]).into_array(), + PrimitiveArray::from_option_iter([Some(4.5f32), None, Some(1.5)]).into_array(), +)] +#[case::f64( + PrimitiveArray::from_iter([1.5f64, 2.5, 3.5, 4.5, 5.5]).into_array(), + PrimitiveArray::from_option_iter([Some(4.5f64), None, Some(1.5)]).into_array(), +)] +#[case::decimal_i8( + DecimalArray::from_iter([10i8, 20, 30, 40, 50], DecimalDType::new(2, 1)).into_array(), + DecimalArray::from_option_iter([Some(40i8), None, Some(10)], DecimalDType::new(2, 1)) + .into_array(), +)] +#[case::decimal_i256( + DecimalArray::from_iter( + [10i128, 20, 30, 40, 50].map(i256::from_i128), + DecimalDType::new(40, 1), + ) + .into_array(), + DecimalArray::from_option_iter( + [Some(i256::from_i128(40)), None, Some(i256::from_i128(10))], + DecimalDType::new(40, 1), + ) + .into_array(), +)] +fn test_take_execute_kernel_copies_every_fixed_width( + #[case] child: crate::ArrayRef, + #[case] expected: crate::ArrayRef, +) -> VortexResult<()> { + let filter = + FilterArray::new(child, Mask::from_iter([true, false, true, true, false])).into_array(); + let codes = PrimitiveArray::new( + buffer![2u64, 100, 0], + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), + ); + let parent = DictArray::try_new(codes.into_array(), filter.clone())?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let result = execute_parent(&filter, &parent, 1, &mut ctx)?; + + assert_arrays_eq!( + result.expect("filter child should execute its take parent"), + expected, + &mut ctx + ); + Ok(()) +} + #[test] fn test_take_execute_kernel_handles_decimal_filter_child() -> VortexResult<()> { let decimal_dtype = DecimalDType::new(19, 2);