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);