diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index 7f7a3f0e441..91700a79d23 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -63,7 +63,7 @@ impl RowFn for NumericBinary { ScalarFnVTable::id(&Binary) } - fn dispatch>( + fn dispatch( &self, op: &Self::Options, args: &[DType], @@ -89,7 +89,7 @@ fn visit_checked(visitor: V) -> VortexResult where T: NativePType, Op: CheckedPrimitiveOp, - V: RowVisitor, + V: RowVisitor, { visitor.visit_deferred::<(T, T), T, Op::Fail>( |(lhs, rhs)| Op::apply(lhs, rhs), @@ -106,7 +106,7 @@ where fn visit_div(visitor: V) -> VortexResult where T: CheckedArithmetic, - V: RowVisitor, + V: RowVisitor, { if T::PTYPE.is_float() { return visit_checked::(visitor); diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index d45a66529a3..3f9d7ae7d3d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -3,22 +3,22 @@ //! Execution arguments paired with the metadata selected during planning. //! -//! [`BorrowedRowFnArgs`] can point at original or sliced arrays while retaining the dtypes, -//! output dtype, and execution policy of the original batch plan. +//! [`BorrowedRowFnArgs`] can point at original or sliced arrays while retaining the dtypes and +//! [`BatchPlan`] of the original batch. use vortex_error::VortexResult; use vortex_error::vortex_err; -use super::RowPolicy; +use super::BatchPlan; use crate::ArrayRef; use crate::dtype::DType; use crate::scalar_fn::ExecutionArgs; /// A borrowed [`ExecutionArgs`] view with the metadata selected for its row function. /// -/// `arrays` can be sliced, while `dtypes` and `output_dtype` always describe the original planned -/// batch. Keeping them together prevents an execution path from pairing an input view with -/// unrelated planning metadata. +/// `arrays` can be sliced, while `dtypes` and `plan` always describe the original planned batch. +/// Keeping them together prevents an execution path from pairing an input view with unrelated +/// planning metadata. #[derive(Clone, Copy)] pub(crate) struct BorrowedRowFnArgs<'a> { /// The input arrays for this row-function invocation. @@ -30,11 +30,8 @@ pub(crate) struct BorrowedRowFnArgs<'a> { /// The original input dtypes used to select the row implementation. dtypes: &'a [DType], - /// The non-nullable dtype built by the selected output capability. - output_dtype: &'a DType, - - /// The nullable execution policy selected during planning. - policy: RowPolicy, + /// The plan an executing dispatch must reproduce. + plan: &'a BatchPlan, } impl<'a> BorrowedRowFnArgs<'a> { @@ -43,15 +40,13 @@ impl<'a> BorrowedRowFnArgs<'a> { arrays: &'a [ArrayRef], row_count: usize, dtypes: &'a [DType], - output_dtype: &'a DType, - policy: RowPolicy, + plan: &'a BatchPlan, ) -> Self { Self { arrays, row_count, dtypes, - output_dtype, - policy, + plan, } } @@ -60,14 +55,9 @@ impl<'a> BorrowedRowFnArgs<'a> { self.dtypes } - /// Return the non-nullable dtype built by the selected output capability. - pub(crate) fn output_dtype(&self) -> &'a DType { - self.output_dtype - } - - /// Return the nullable execution policy selected during planning. - pub(crate) fn policy(&self) -> RowPolicy { - self.policy + /// Return the plan an executing dispatch must reproduce. + pub(crate) fn plan(&self) -> &'a BatchPlan { + self.plan } } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs index 68a235de466..32cdee5ed0a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -73,7 +73,7 @@ impl RowFnExecutionArgs { return self.execute_dense(kernel, ctx); } - match self.policy { + match self.plan.policy() { RowPolicy::Dense => self.execute_dense(kernel, ctx), RowPolicy::DenseWithRetry => { self.execute_dense_with_retry(execute_dense_attempt, try_valid_rows, ctx) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs index 0f8170630e3..3a83fba60e1 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs @@ -20,25 +20,37 @@ impl RowFnExecutionArgs { ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() } - /// Validate the finished output and apply the row function's logical outer nullability. + /// Label the finished column, validate it, and apply the row function's logical outer + /// nullability. pub(super) fn finalize_output( &self, values: ArrayRef, expected_len: usize, ) -> VortexResult { + // Label before validation so the checks below see the dtype this function returns. Every + // batch strategy reaches this method after masking, so an empty, all-null, constant, or + // partially valid batch carries the same metadata as a dense one. + let values = self.plan.relabel_output(values)?; + validate_output(self.id, &self.result_dtype, expected_len, &values)?; cast_output_nullability(&self.result_dtype, values) } - /// Validate the output from a row function before batch validity is attached. + /// Validate the unlabelled output from a row function before batch validity is attached. pub(super) fn validate_kernel_output( &self, values: ArrayRef, expected_len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - finalize_kernel_output(self.id, &self.output_dtype, expected_len, values, ctx) + finalize_kernel_output( + self.id, + self.plan.storage_dtype(), + expected_len, + values, + ctx, + ) } } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs index b2784009011..36b6c246a27 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -49,15 +49,12 @@ pub(crate) struct RowFnExecutionArgs { /// input. Conjoining is lazy, and null handling materializes the mask only when required. validity: Validity, - /// The declared output dtype, widened to nullable when any input is nullable. Kernel output is + /// The output dtype, widened to nullable when any input is nullable. The finished column is /// reconciled against this dtype. result_dtype: DType, - /// The non-nullable dtype the dispatched output capability builds, computed while planning. - output_dtype: DType, - - /// How the concrete dispatch executes nullable rows. - policy: RowPolicy, + /// The storage dtype, output label, and null-handling policy selected while planning. + plan: BatchPlan, } #[cfg(test)] diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs index fa52fd91b05..0b7eb0aa728 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs @@ -55,8 +55,7 @@ impl RowFnExecutionArgs { arg_dtypes, validity, result_dtype, - output_dtype: plan.output_dtype, - policy: plan.policy, + plan, }) } @@ -66,12 +65,6 @@ impl RowFnExecutionArgs { arrays: &'b [ArrayRef], row_count: usize, ) -> BorrowedRowFnArgs<'b> { - BorrowedRowFnArgs::new( - arrays, - row_count, - &self.arg_dtypes, - &self.output_dtype, - self.policy, - ) + BorrowedRowFnArgs::new(arrays, row_count, &self.arg_dtypes, &self.plan) } } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 78dcb554a07..224c875c52b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use rstest::rstest; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -16,11 +17,15 @@ use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::ConstantArray; +use crate::arrays::ExtensionArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; +use crate::dtype::extension::ExtDTypeRef; +use crate::extension::datetime::TimeUnit; +use crate::extension::datetime::Timestamp; use crate::scalar::Scalar; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; @@ -30,6 +35,7 @@ use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::RowVisitor; use crate::scalar_fn::unstable::row::execute_rows; +use crate::scalar_fn::unstable::row::row_fn_return_dtype; use crate::validity::Validity; #[derive(Clone, Default)] @@ -71,13 +77,13 @@ struct I64Sink(BufferMut); // SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that // initialized slice. The `()` write token therefore proves no additional invariant. -unsafe impl OutputSink for I64Sink { +unsafe impl OutputSink for I64Sink { type Rows<'a> = &'a mut [i64]; type Row<'a> = &'a mut i64; type WriteToken = (); - fn return_dtype(_options: &Options) -> VortexResult { - Ok(DType::from(i64::PTYPE)) + fn storage_dtype() -> DType { + DType::from(i64::PTYPE) } fn with_capacity(rows: usize) -> VortexResult { @@ -109,7 +115,7 @@ impl RowFn for DeferredAdd { *ID } - fn dispatch>( + fn dispatch( &self, _options: &Self::Options, _args: &[DType], @@ -144,7 +150,7 @@ impl RowFn for ValidOnlyIdentity { *ID } - fn dispatch>( + fn dispatch( &self, _options: &Self::Options, _args: &[DType], @@ -168,7 +174,7 @@ impl RowFn for InvalidKernelOutput { *ID } - fn dispatch>( + fn dispatch( &self, _options: &Self::Options, _args: &[DType], @@ -312,3 +318,335 @@ fn test_valid_only_empty_batch_preserves_nonnullable_dtype() -> VortexResult<()> assert_eq!(actual.dtype(), &DType::from(i64::PTYPE)); Ok(()) } + +/// A timestamp extension dtype over `i64` storage of the given nullability. +fn timestamp_ext(nullability: Nullability) -> ExtDTypeRef { + Timestamp::new(TimeUnit::Seconds, nullability).erased() +} + +/// The output dtype a declaring dispatch labels onto its `i64` storage. +fn timestamp_dtype(nullability: Nullability) -> DType { + DType::Extension(timestamp_ext(nullability)) +} + +/// Build the timestamp column a labelled dispatch is expected to produce. +fn expected_timestamps(values: Vec, validity: Validity) -> VortexResult { + let storage = PrimitiveArray::new(values, validity).into_array(); + let ext_dtype = timestamp_ext(storage.dtype().nullability()); + + Ok(ExtensionArray::try_new(ext_dtype, storage)?.into_array()) +} + +/// Returns its input unchanged under the output dtype each dispatch declares. +/// +/// `declared` is `None` to leave the storage dtype in place, so one function covers both the +/// labelled and unlabelled paths and every rejected label. +#[derive(Clone)] +struct DeclaredOutput { + declared: Option, +} + +impl DeclaredOutput { + fn timestamps() -> Self { + Self { + declared: Some(timestamp_dtype(Nullability::NonNullable)), + } + } +} + +impl RowFn for DeclaredOutput { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.declared_output"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + let visitor = match &self.declared { + Some(declared) => visitor.with_output_dtype(declared.clone()), + None => visitor, + }; + + visitor.visit::<(i64,), i64>(|(value,)| value) + } +} + +/// Labels the output of a sink-writing dispatch, which plans [`RowPolicy::ValidOnly`]. +#[derive(Clone)] +struct DeclaredSinkOutput; + +impl RowFn for DeclaredSinkOutput { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.declared_sink_output"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor + .with_output_dtype(timestamp_dtype(Nullability::NonNullable)) + .visit_into::<(i64,), I64Sink, VortexResult<()>>(|(value,), output| { + *output = value; + Ok(()) + }) + } +} + +/// Declares a different output dtype on its second dispatch, which execution must reject. +#[derive(Clone)] +struct ChangingOutputDType { + dispatches: Arc, +} + +impl RowFn for ChangingOutputDType { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.changing_output_dtype"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + let visitor = if self.dispatches.fetch_add(1, Ordering::Relaxed) == 0 { + visitor.with_output_dtype(timestamp_dtype(Nullability::NonNullable)) + } else { + visitor + }; + + visitor.visit::<(i64,), i64>(|(value,)| value) + } +} + +#[rstest] +#[case::all_valid(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false, true]))] +#[case::array_backed_all_valid(Validity::Array(ConstantArray::new(true, 3).into_array()))] +fn test_declared_output_dtype_labels_batch(#[case] validity: Validity) -> VortexResult<()> { + let values = vec![1_i64, 2, 3]; + let input = PrimitiveArray::new(values.clone(), validity.clone()).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows( + &DeclaredOutput::timestamps(), + &EmptyOptions, + &args, + &mut ctx, + )?; + let expected = expected_timestamps(values, validity)?; + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_declared_output_dtype_labels_sink_output() -> VortexResult<()> { + // `I64Sink` declines skip-invalid execution, so this batch stays all-valid. The masked path is + // covered by the owned-output cases above. + let values = vec![1_i64, 2, 3]; + let input = PrimitiveArray::new(values.clone(), Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&DeclaredSinkOutput, &EmptyOptions, &args, &mut ctx)?; + let expected = expected_timestamps(values, Validity::NonNullable)?; + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_declared_output_dtype_labels_empty_batch() -> VortexResult<()> { + let input = PrimitiveArray::from_iter(std::iter::empty::()).into_array(); + let args = VecExecutionArgs::new(vec![input], 0); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows( + &DeclaredOutput::timestamps(), + &EmptyOptions, + &args, + &mut ctx, + )?; + + assert_eq!(actual.len(), 0); + assert_eq!(actual.dtype(), ×tamp_dtype(Nullability::NonNullable)); + Ok(()) +} + +#[test] +fn test_declared_output_dtype_labels_all_null_batch() -> VortexResult<()> { + let null_i64 = Scalar::null(DType::Primitive(i64::PTYPE, Nullability::Nullable)); + let input = ConstantArray::new(null_i64, 3).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows( + &DeclaredOutput::timestamps(), + &EmptyOptions, + &args, + &mut ctx, + )?; + let expected = + ConstantArray::new(Scalar::null(timestamp_dtype(Nullability::Nullable)), 3).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_declared_output_dtype_labels_constant_batch() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 3).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows( + &DeclaredOutput::timestamps(), + &EmptyOptions, + &args, + &mut ctx, + )?; + let expected = expected_timestamps(vec![7, 7, 7], Validity::NonNullable)?; + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_declared_output_dtype_reaches_planning() -> VortexResult<()> { + let args = [DType::Primitive(i64::PTYPE, Nullability::Nullable)]; + + let dtype = row_fn_return_dtype(&DeclaredOutput::timestamps(), &EmptyOptions, &args)?; + + assert_eq!(dtype, timestamp_dtype(Nullability::Nullable)); + Ok(()) +} + +#[rstest] +#[case::nullable(timestamp_dtype(Nullability::Nullable), "must be non-nullable")] +#[case::not_an_extension(DType::from(u64::PTYPE), "must be an extension dtype")] +fn test_declared_output_dtype_rejects_bad_label( + #[case] declared: DType, + #[case] expected_message: &str, +) -> VortexResult<()> { + let function = DeclaredOutput { + declared: Some(declared), + }; + let input = PrimitiveArray::from_iter(vec![1_i64, 2]).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = execute_rows(&function, &EmptyOptions, &args, &mut ctx) + .expect_err("an invalid output dtype must be rejected") + .to_string(); + + assert!( + error.contains(expected_message), + "the label error must contain {expected_message:?}, got {error}", + ); + Ok(()) +} + +/// Declares a timestamp output dtype over `u64` storage, which its `i64` storage cannot match. +#[derive(Clone)] +struct MismatchedStorage; + +impl RowFn for MismatchedStorage { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.mismatched_storage"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor + .with_output_dtype(timestamp_dtype(Nullability::NonNullable)) + .visit::<(u64,), u64>(|(value,)| value) + } +} + +#[test] +fn test_declared_output_dtype_rejects_mismatched_storage() -> VortexResult<()> { + let input = PrimitiveArray::from_iter(vec![1_u64, 2]).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = execute_rows(&MismatchedStorage, &EmptyOptions, &args, &mut ctx) + .expect_err("an extension label over another storage dtype must be rejected") + .to_string(); + + assert!( + error.contains("must store"), + "the label error must report the expected storage dtype, got {error}", + ); + Ok(()) +} + +#[test] +fn test_execution_rejects_a_changed_output_dtype() -> VortexResult<()> { + let function = ChangingOutputDType { + dispatches: Arc::new(AtomicUsize::new(0)), + }; + let input = PrimitiveArray::from_iter(vec![1_i64, 2]).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = execute_rows(&function, &EmptyOptions, &args, &mut ctx) + .expect_err("an execution dispatch must declare the planned output dtype") + .to_string(); + + assert!( + error.contains("must declare the planned output dtype"), + "execution must reject a changed output dtype, got {error}", + ); + Ok(()) +} + +#[test] +fn test_undeclared_output_dtype_keeps_the_storage_dtype() -> VortexResult<()> { + let function = DeclaredOutput { declared: None }; + let values = vec![1_i64, 2]; + let input = PrimitiveArray::from_iter(values.clone()).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&function, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter(values).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs index 9903d53530a..7ed012191ea 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -26,16 +26,16 @@ use crate::scalar_fn::unstable::row::ViewLen; /// The executor owns the sink and passes each output row to `apply`. This keeps `apply` as [`Fn`]. /// Capturing the sink would require [`FnMut`] and put its buffer metadata behind loop-carried /// mutable closure state, which can prevent LLVM from treating that metadata as loop-invariant. -pub(crate) fn execute_sink( +pub(crate) fn execute_sink( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult>::WriteToken>, + Sink: OutputSink, + ApplyResult: SinkResult, { let columns = Args::decode(args, ctx)?; @@ -43,11 +43,11 @@ where let const_values = Args::const_values(&columns); let prepared = prepare(const_values); - let mut sink = >::with_capacity(row_count)?; + let mut sink = Sink::with_capacity(row_count)?; // Keep `rows` scoped so its borrow ends before `finish`, which consumes the sink. { - let mut rows = >::rows(&mut sink); + let mut rows = Sink::rows(&mut sink); // This equality proves to LLVM that `0..row_count` is in bounds for `rows`. let sink_row_count = rows.len(); @@ -68,8 +68,7 @@ where // `row_count` rows before the loop. let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; // SAFETY: the sink row-count check above proved every loop index is in bounds. - let output = - unsafe { >::row_unchecked(&mut rows, index) }; + let output = unsafe { Sink::row_unchecked(&mut rows, index) }; apply(&prepared, elements, output).into_result()?; } @@ -80,8 +79,7 @@ where for index in 0..row_count { // SAFETY: the sink row-count check above proved every loop index is in bounds. - let output = - unsafe { >::row_unchecked(&mut rows, index) }; + let output = unsafe { Sink::row_unchecked(&mut rows, index) }; // LLVM unswitches the batch-constant checks in `Args::get` before vectorizing the // loop. @@ -91,7 +89,7 @@ where } // SAFETY: every row callback completed successfully, so each returned the required write token. - unsafe { >::finish(sink) } + unsafe { Sink::finish(sink) } } /// Write only the rows set in `valid`, or decline when the inputs or sink cannot support @@ -99,17 +97,17 @@ where /// /// `Ok(None)` signals that direct skip-invalid execution is unavailable. Batch execution decides /// how to handle the decline. -pub(crate) fn execute_sink_valid_rows( +pub(crate) fn execute_sink_valid_rows( args: &dyn ExecutionArgs, valid: &MaskValuesRef, ctx: &mut ExecutionCtx, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult> where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult>::WriteToken>, + Sink: OutputSink, + ApplyResult: SinkResult, { let Some(ValidRowsSetup { initialize_skipped_rows, @@ -117,7 +115,7 @@ where valid_rows, row_count, mut sink, - }) = setup_sink_valid_rows::(args, valid, ctx)? + }) = setup_sink_valid_rows::(args, valid, ctx)? else { return Ok(None); }; @@ -130,7 +128,7 @@ where // `drop(rows)` duplicates `Args::get` in every sparse callback. { // Initialize every slot before visiting only valid rows. - let mut rows = >::rows(&mut sink); + let mut rows = Sink::rows(&mut sink); initialize_skipped_rows(&mut rows); // The initializer can change addressability. Recheck it so LLVM can prove every mask @@ -150,8 +148,7 @@ where valid_rows.try_for_each_set_index(|index| { // SAFETY: the post-initialization row-count check proved that the sink addresses // every mask index, which is below the mask's validated `row_count`. - let output = - unsafe { >::row_unchecked(&mut rows, index) }; + let output = unsafe { Sink::row_unchecked(&mut rows, index) }; // SAFETY: `view_lens_match` checked that these exact retained views address // `row_count` rows, and mask indices are below `row_count`. @@ -167,8 +164,7 @@ where valid_rows.try_for_each_set_index(|index| { // SAFETY: the post-initialization row-count check proved that the sink addresses // every mask index, which is below the mask's validated `row_count`. - let output = - unsafe { >::row_unchecked(&mut rows, index) }; + let output = unsafe { Sink::row_unchecked(&mut rows, index) }; apply(&prepared, Args::get(&columns, index), output).into_result() })?; @@ -177,7 +173,7 @@ where // SAFETY: the initializer completed before traversal, and every visited callback completed // successfully and returned the required write token. - unsafe { >::finish(sink) }.map(Some) + unsafe { Sink::finish(sink) }.map(Some) } /// Construct a decoded-length error outside the traversal branches. @@ -193,12 +189,12 @@ fn decoded_length_error(row_count: usize) -> VortexResult<()> { } /// State resolved before preparing the skip-invalid row loop. -struct ValidRowsSetup<'valid, Args, Sink, Options> +struct ValidRowsSetup<'valid, Args, Sink> where Args: ElementTuple, - Sink: OutputSink, + Sink: OutputSink, { - initialize_skipped_rows: for<'rows> fn(&mut >::Rows<'rows>), + initialize_skipped_rows: for<'rows> fn(&mut Sink::Rows<'rows>), columns: Args::Columns, valid_rows: &'valid BitBuffer, row_count: usize, @@ -206,18 +202,17 @@ where } /// Resolve the capabilities, inputs, sink, and validity mask for skip-invalid execution. -fn setup_sink_valid_rows<'valid, Args, Sink, Options>( +fn setup_sink_valid_rows<'valid, Args, Sink>( args: &dyn ExecutionArgs, valid: &'valid MaskValuesRef, ctx: &mut ExecutionCtx, -) -> VortexResult>> +) -> VortexResult>> where Args: ElementTuple, - Sink: OutputSink, + Sink: OutputSink, { // The initializer both declares support for skipping rows and initializes those rows. - let Some(initialize_skipped_rows) = >::skipped_rows_initializer() - else { + let Some(initialize_skipped_rows) = Sink::skipped_rows_initializer() else { return Ok(None); }; @@ -232,7 +227,7 @@ where // Keep allocation before the validity and length checks. With multiple CGUs and no LTO, // moving it later inlines `Args::get` into every sparse callback, duplicating its bounds // checks. - let sink = >::with_capacity(row_count)?; + let sink = Sink::with_capacity(row_count)?; let valid_rows = valid.bit_buffer(); vortex_ensure_eq!( @@ -266,7 +261,6 @@ mod tests { use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::NativePType; - use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::OutputSink; use crate::validity::Validity; @@ -277,13 +271,13 @@ mod tests { // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or // `finish` through the executor. The row-initialization requirements are therefore vacuous. - unsafe impl OutputSink for NonSkippingSink { + unsafe impl OutputSink for NonSkippingSink { type Rows<'a> = (); type Row<'a> = (); type WriteToken = (); - fn return_dtype(_options: &Options) -> VortexResult { - Ok(DType::from(i64::PTYPE)) + fn storage_dtype() -> DType { + DType::from(i64::PTYPE) } fn with_capacity(_rows: usize) -> VortexResult { @@ -305,7 +299,7 @@ mod tests { // SAFETY: the initializer deliberately shrinks the row collection to exercise the executor's // post-initialization length check. If execution incorrectly continues, safe indexing in // `row_unchecked` panics instead of accessing invalid memory. - unsafe impl OutputSink for ShrinkingSink { + unsafe impl OutputSink for ShrinkingSink { type Rows<'a> = &'a mut Vec; type Row<'a> = &'a mut i64; type WriteToken = (); @@ -316,8 +310,8 @@ mod tests { }) } - fn return_dtype(_options: &Options) -> VortexResult { - Ok(DType::from(i64::PTYPE)) + fn storage_dtype() -> DType { + DType::from(i64::PTYPE) } fn with_capacity(rows: usize) -> VortexResult { @@ -346,7 +340,7 @@ mod tests { }; let mut ctx = array_session().create_execution_ctx(); - let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, (), EmptyOptions>( + let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, ()>( &args, &valid, &mut ctx, @@ -368,7 +362,7 @@ mod tests { }; let mut ctx = array_session().create_execution_ctx(); - let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, (), EmptyOptions>( + let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, ()>( &args, &valid, &mut ctx, diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index e456886e618..55087124b72 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -77,7 +77,7 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { /// /// Planning and execution both call this method, so its result **must** depend only on /// `options` and `args`. Cross-argument dtype validation belongs here. - fn dispatch>( + fn dispatch( &self, options: &Self::Options, args: &[DType], diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs index 2bf51cf999d..40f449ff4ae 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -18,8 +18,13 @@ use crate::scalar_fn::unstable::row::ViewLen; /// A column allocated once per batch that a row closure writes into, one row at a time. /// -/// A sink can use function options to build a runtime-shaped output or own shared batch state. The -/// executor passes each row slot into an [`Fn`] closure. +/// A sink owns batch-wide state that an independent owned value cannot express, such as +/// uninitialized storage or a row handle covering more than one element. The executor passes each +/// row slot into an [`Fn`] closure. +/// +/// A sink describes only how rows are physically written. An output dtype derived from the +/// function options or argument dtypes is declared by [`RowVisitor::with_output_dtype`], which +/// labels the column this sink builds. /// /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once. /// Execution can omit invalid rows when [`skipped_rows_initializer`] returns an initializer. @@ -55,9 +60,10 @@ use crate::scalar_fn::unstable::row::ViewLen; /// [`WriteToken`]: Self::WriteToken /// [`finish`]: Self::finish /// [`RowFn::INFALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::INFALLIBLE +/// [`RowVisitor::with_output_dtype`]: crate::scalar_fn::unstable::row::RowVisitor::with_output_dtype /// [`SinkResult`]: crate::scalar_fn::unstable::row::SinkResult /// [`skipped_rows_initializer`]: Self::skipped_rows_initializer -pub unsafe trait OutputSink: 'static + Sized { +pub unsafe trait OutputSink: 'static + Sized { /// A loop-local view of all output rows. /// /// Borrowed once before execution so the sink's buffer descriptor and shape become loop @@ -89,11 +95,15 @@ pub unsafe trait OutputSink: 'static + Sized { None } - /// The dtype of the column this sink builds, given the function options. + /// The dtype of the column this sink builds. + /// + /// Because this method takes no arguments, the dtype must be a property of the Rust type. An + /// output dtype that depends on the function options or argument dtypes is declared by + /// [`RowVisitor::with_output_dtype`](crate::scalar_fn::unstable::row::RowVisitor::with_output_dtype). /// /// **Must** be non-nullable: batch execution derives nullability from the inputs, widens the /// result, and masks the null rows. - fn return_dtype(options: &Options) -> VortexResult; + fn storage_dtype() -> DType; /// Allocate a sink for `rows` rows. fn with_capacity(rows: usize) -> VortexResult; @@ -109,7 +119,7 @@ pub unsafe trait OutputSink: 'static + Sized { unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; /// Finish into the built column, whose dtype **must** be this sink's - /// [`return_dtype`](Self::return_dtype). Called once per batch. + /// [`storage_dtype`](Self::storage_dtype). Called once per batch. /// /// # Safety /// @@ -171,9 +181,7 @@ pub struct UninitElementSink { // names one distinct slot. Safe code cannot construct `InitializedElement`. Its unsafe constructor // writes the supplied slot and requires the caller to return that exact evidence. The // skipped-row initializer writes `T::default()` into every slot before masked traversal. -unsafe impl OutputSink - for UninitElementSink -{ +unsafe impl OutputSink for UninitElementSink { type Rows<'a> = &'a mut [MaybeUninit]; type Row<'a> = &'a mut MaybeUninit; type WriteToken = InitializedElement; @@ -186,8 +194,8 @@ unsafe impl OutputSink }) } - fn return_dtype(_options: &Options) -> VortexResult { - Ok(T::element_dtype()) + fn storage_dtype() -> DType { + T::element_dtype() } fn with_capacity(rows: usize) -> VortexResult { diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs index 5611e779b44..cf5bf7ad44f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs @@ -96,17 +96,14 @@ pub(super) fn validate_owned_visit( Ok(dtype) } -pub(super) fn validate_sink_visit( - options: &Options, - dtypes: &[DType], -) -> VortexResult +pub(super) fn validate_sink_visit(dtypes: &[DType]) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, + Sink: OutputSink, { Args::validate(dtypes)?; - let dtype = Sink::return_dtype(options)?; + let dtype = Sink::storage_dtype(); vortex_ensure!( !dtype.is_nullable(), "row output sinks must declare a non-nullable dtype, got {dtype}", diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index 8fc7209bc19..435ae9698d9 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -3,13 +3,16 @@ //! Visitors that execute dense and skip-invalid row loops. //! -//! Each visit revalidates its concrete signature and checks that its output dtype and execution -//! policy match the plan before entering a row loop. [`ExecuteValidRows`] can decline when the -//! signature cannot execute over the original inputs. +//! Each visit revalidates its concrete signature and checks that its plan matches the one planning +//! selected before entering a row loop. [`ExecuteValidRows`] can decline when the signature cannot +//! execute over the original inputs. + +use std::marker::PhantomData; use vortex_error::VortexResult; use vortex_mask::MaskValuesRef; +use super::BatchPlan; use super::RowPolicy; use super::RowVisitor; use super::check::assert_deferred_visit_contract; @@ -17,7 +20,6 @@ use super::check::assert_owned_visit_contract; use super::check::assert_sink_visit_contract; use super::check::validate_owned_visit; use super::check::validate_sink_visit; -use super::ensure_plan; use super::row_visitor::private; use crate::ArrayRef; use crate::ExecutionCtx; @@ -45,44 +47,47 @@ pub(crate) struct ExecuteRows<'args, 'ctx, F: RowFn> { /// The input dtypes used by the planning visit. dtypes: &'args [DType], - /// The function options used to derive a sink's runtime dtype. - options: &'args F::Options, - - /// The output dtype computed by the planning visit. - output_dtype: &'args DType, + /// The plan selected by the planning visit, which this visit must reproduce. + plan: &'args BatchPlan, - /// The nullable execution policy selected by the planning visit. - policy: RowPolicy, + /// The output dtype declared by [`RowVisitor::with_output_dtype`], if any. + output_dtype: Option, /// The execution context used to decode the input columns. ctx: &'ctx mut ExecutionCtx, + + /// Ties this visit to the function used by its compile-time contract checks. + function: PhantomData, } impl<'args, 'ctx, F: RowFn> ExecuteRows<'args, 'ctx, F> { pub(crate) fn new( args: &'args dyn ExecutionArgs, dtypes: &'args [DType], - options: &'args F::Options, - output_dtype: &'args DType, - policy: RowPolicy, + plan: &'args BatchPlan, ctx: &'ctx mut ExecutionCtx, ) -> Self { Self { args, dtypes, - options, - output_dtype, - policy, + plan, + output_dtype: None, ctx, + function: PhantomData, } } } impl private::Sealed for ExecuteRows<'_, '_, F> {} -impl RowVisitor for ExecuteRows<'_, '_, F> { +impl RowVisitor for ExecuteRows<'_, '_, F> { type VisitResult = ArrayRef; + fn with_output_dtype(mut self, dtype: DType) -> Self { + self.output_dtype = Some(dtype); + self + } + fn visit_prepared( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, @@ -93,12 +98,12 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { Out: OutputElement, { const { assert_owned_visit_contract::() }; - ensure_plan( - self.output_dtype, - self.policy, + let visited = BatchPlan::new( validate_owned_visit::(self.dtypes)?, + self.output_dtype, RowPolicy::for_owned_output::(), )?; + self.plan.ensure_reproduced_by(&visited)?; execute_owned_infallible::(self.args, self.ctx, prepare, apply) } @@ -106,28 +111,22 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { fn visit_prepared_into( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - apply: impl Fn( - &Prepared, - Args::Elems<'_>, - >::Row<'_>, - ) -> ApplyResult, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult>::WriteToken>, + Sink: OutputSink, + ApplyResult: SinkResult, { const { assert_sink_visit_contract::() }; - ensure_plan( + let visited = BatchPlan::new( + validate_sink_visit::(self.dtypes)?, self.output_dtype, - self.policy, - validate_sink_visit::(self.options, self.dtypes)?, RowPolicy::for_sink::(), )?; + self.plan.ensure_reproduced_by(&visited)?; - execute_sink::( - self.args, self.ctx, prepare, apply, - ) + execute_sink::(self.args, self.ctx, prepare, apply) } fn visit_prepared_deferred( @@ -142,12 +141,12 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { Fail: FailureEvidence, { const { assert_deferred_visit_contract::() }; - ensure_plan( - self.output_dtype, - self.policy, + let visited = BatchPlan::new( validate_owned_visit::(self.dtypes)?, + self.output_dtype, RowPolicy::for_deferred_output::(), )?; + self.plan.ensure_reproduced_by(&visited)?; execute_owned::( self.args, @@ -170,49 +169,52 @@ pub(crate) struct ExecuteValidRows<'args, 'ctx, F: RowFn> { /// The input dtypes used by the planning visit. dtypes: &'args [DType], - /// The function options used to derive a sink's runtime dtype. - options: &'args F::Options, - - /// The output dtype computed by the planning visit. - output_dtype: &'args DType, + /// The plan selected by the planning visit, which this visit must reproduce. + plan: &'args BatchPlan, - /// The nullable execution policy selected by the planning visit. - policy: RowPolicy, + /// The output dtype declared by [`RowVisitor::with_output_dtype`], if any. + output_dtype: Option, /// The conjoined validity, containing both valid and invalid rows. valid: MaskValuesRef, /// The execution context used to decode the input columns. ctx: &'ctx mut ExecutionCtx, + + /// Ties this visit to the function used by its compile-time contract checks. + function: PhantomData, } impl<'args, 'ctx, F: RowFn> ExecuteValidRows<'args, 'ctx, F> { pub(crate) fn new( args: &'args dyn ExecutionArgs, dtypes: &'args [DType], - options: &'args F::Options, - output_dtype: &'args DType, - policy: RowPolicy, + plan: &'args BatchPlan, valid: MaskValuesRef, ctx: &'ctx mut ExecutionCtx, ) -> Self { Self { args, dtypes, - options, - output_dtype, - policy, + plan, + output_dtype: None, valid, ctx, + function: PhantomData, } } } impl private::Sealed for ExecuteValidRows<'_, '_, F> {} -impl RowVisitor for ExecuteValidRows<'_, '_, F> { +impl RowVisitor for ExecuteValidRows<'_, '_, F> { type VisitResult = Option; + fn with_output_dtype(mut self, dtype: DType) -> Self { + self.output_dtype = Some(dtype); + self + } + fn visit_prepared( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, @@ -223,12 +225,12 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { Out: OutputElement, { const { assert_owned_visit_contract::() }; - ensure_plan( - self.output_dtype, - self.policy, + let visited = BatchPlan::new( validate_owned_visit::(self.dtypes)?, + self.output_dtype, RowPolicy::for_owned_output::(), )?; + self.plan.ensure_reproduced_by(&visited)?; execute_owned_infallible_valid_rows::( self.args, @@ -242,26 +244,22 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { fn visit_prepared_into( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - apply: impl Fn( - &Prepared, - Args::Elems<'_>, - >::Row<'_>, - ) -> ApplyResult, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult>::WriteToken>, + Sink: OutputSink, + ApplyResult: SinkResult, { const { assert_sink_visit_contract::() }; - ensure_plan( + let visited = BatchPlan::new( + validate_sink_visit::(self.dtypes)?, self.output_dtype, - self.policy, - validate_sink_visit::(self.options, self.dtypes)?, RowPolicy::for_sink::(), )?; + self.plan.ensure_reproduced_by(&visited)?; - execute_sink_valid_rows::( + execute_sink_valid_rows::( self.args, &self.valid, self.ctx, @@ -282,12 +280,12 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { Fail: FailureEvidence, { const { assert_deferred_visit_contract::() }; - ensure_plan( - self.output_dtype, - self.policy, + let visited = BatchPlan::new( validate_owned_visit::(self.dtypes)?, + self.output_dtype, RowPolicy::for_deferred_output::(), )?; + self.plan.ensure_reproduced_by(&visited)?; execute_owned_valid_rows::( self.args, diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs index 8c0bfea0027..2c90d8a5a08 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -5,11 +5,6 @@ //! //! [`RowFn::dispatch`]: crate::scalar_fn::unstable::row::RowFn::dispatch -use vortex_error::VortexResult; -use vortex_error::vortex_ensure_eq; - -use crate::dtype::DType; - mod check; pub(super) use check::assert_owned_output_needs_no_drop; @@ -27,23 +22,3 @@ pub(super) use plan::RowPolicy; mod row_visitor; pub use row_visitor::RowVisitor; - -fn ensure_plan( - planned_output: &DType, - planned_policy: RowPolicy, - actual_output: DType, - actual_policy: RowPolicy, -) -> VortexResult<()> { - vortex_ensure_eq!( - actual_policy, - planned_policy, - "row dispatch must select the planned nullable execution policy: planned {planned_policy:?}, got {actual_policy:?}", - ); - vortex_ensure_eq!( - actual_output, - *planned_output, - "row dispatch must select the planned output dtype: planned {planned_output}, got {actual_output}", - ); - - Ok(()) -} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs index 9299c1d59fc..f603ca6e685 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -3,12 +3,16 @@ //! Plans the concrete signature selected by [`RowFn::dispatch`]. //! -//! [`BatchPlanner`] validates input and output dtypes, then records the output dtype and -//! null-handling policy that execution must reproduce. +//! [`BatchPlanner`] validates input and output dtypes, then records the [`BatchPlan`] that +//! execution must reproduce: the dtype the dispatched capability builds, the dtype the function +//! returns, and the null-handling policy. use std::marker::PhantomData; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use super::RowVisitor; use super::check::assert_deferred_visit_contract; @@ -17,8 +21,12 @@ use super::check::assert_sink_visit_contract; use super::check::validate_owned_visit; use super::check::validate_sink_visit; use super::row_visitor::private; +use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::ExtensionArray; use crate::dtype::DType; use crate::dtype::Nullability; +use crate::dtype::extension::ExtDTypeRef; use crate::scalar_fn::unstable::row::ElementTuple; use crate::scalar_fn::unstable::row::FailureEvidence; use crate::scalar_fn::unstable::row::IndexedElementTuple; @@ -31,17 +39,18 @@ use crate::scalar_fn::unstable::row::SinkResult; pub(crate) struct BatchPlanner<'a, F: RowFn> { dtypes: &'a [DType], - options: &'a F::Options, + /// The output dtype declared by [`RowVisitor::with_output_dtype`], if any. + output_dtype: Option, /// Ties the planner to the function used by its compile-time contract checks. function: PhantomData, } impl<'a, F: RowFn> BatchPlanner<'a, F> { - pub(crate) fn new(dtypes: &'a [DType], options: &'a F::Options) -> Self { + pub(crate) fn new(dtypes: &'a [DType]) -> Self { Self { dtypes, - options, + output_dtype: None, function: PhantomData, } } @@ -49,9 +58,14 @@ impl<'a, F: RowFn> BatchPlanner<'a, F> { impl private::Sealed for BatchPlanner<'_, F> {} -impl RowVisitor for BatchPlanner<'_, F> { +impl RowVisitor for BatchPlanner<'_, F> { type VisitResult = BatchPlan; + fn with_output_dtype(mut self, dtype: DType) -> Self { + self.output_dtype = Some(dtype); + self + } + fn visit_prepared( self, _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, @@ -62,31 +76,31 @@ impl RowVisitor for BatchPlanner<'_, F> { Out: OutputElement, { const { assert_owned_visit_contract::() }; - Ok(BatchPlan { - output_dtype: validate_owned_visit::(self.dtypes)?, - policy: RowPolicy::for_owned_output::(), - }) + + BatchPlan::new( + validate_owned_visit::(self.dtypes)?, + self.output_dtype, + RowPolicy::for_owned_output::(), + ) } fn visit_prepared_into( self, _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - _apply: impl Fn( - &Prepared, - Args::Elems<'_>, - >::Row<'_>, - ) -> ApplyResult, + _apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult>::WriteToken>, + Sink: OutputSink, + ApplyResult: SinkResult, { const { assert_sink_visit_contract::() }; - Ok(BatchPlan { - output_dtype: validate_sink_visit::(self.options, self.dtypes)?, - policy: RowPolicy::for_sink::(), - }) + + BatchPlan::new( + validate_sink_visit::(self.dtypes)?, + self.output_dtype, + RowPolicy::for_sink::(), + ) } fn visit_prepared_deferred( @@ -101,30 +115,163 @@ impl RowVisitor for BatchPlanner<'_, F> { Fail: FailureEvidence, { const { assert_deferred_visit_contract::() }; - Ok(BatchPlan { - output_dtype: validate_owned_visit::(self.dtypes)?, - policy: RowPolicy::for_deferred_output::(), - }) + + BatchPlan::new( + validate_owned_visit::(self.dtypes)?, + self.output_dtype, + RowPolicy::for_deferred_output::(), + ) } } -/// The execution policy and output dtype selected by a planning visit. +/// The output dtypes and execution policy selected by one planning visit. +/// +/// The _storage dtype_ is what the dispatched [`OutputElement`] or [`OutputSink`] physically +/// builds. The _output dtype_ is what the function returns. They differ when a dispatch declares a +/// label through [`RowVisitor::with_output_dtype`], which is how a function derives an extension +/// output dtype from its options or argument dtypes. pub(crate) struct BatchPlan { - /// The non-nullable dtype built by the selected output capability. - pub(crate) output_dtype: DType, + /// The non-nullable dtype the dispatched output capability builds. Kernel output is validated + /// against this dtype before any label is applied. + storage_dtype: DType, + + /// The extension dtype labelled onto the finished column, when the declared output dtype + /// differs from the storage dtype. + output_label: Option, /// How this concrete dispatch executes nullable rows. - pub(crate) policy: RowPolicy, + policy: RowPolicy, } impl BatchPlan { + /// Plan an output built as `storage_dtype` and returned as `output_dtype`. + /// + /// `output_dtype` is the dtype a dispatch declared through + /// [`RowVisitor::with_output_dtype`], or `None` to return the storage dtype unchanged. + pub(crate) fn new( + storage_dtype: DType, + output_dtype: Option, + policy: RowPolicy, + ) -> VortexResult { + let output_label = match output_dtype { + Some(output_dtype) => validate_output_label(&storage_dtype, output_dtype)?, + None => None, + }; + + Ok(Self { + storage_dtype, + output_label, + policy, + }) + } + + /// Return the non-nullable dtype the dispatched output capability builds. + pub(crate) fn storage_dtype(&self) -> &DType { + &self.storage_dtype + } + + /// Return the non-nullable dtype the function returns. + pub(crate) fn output_dtype(&self) -> DType { + match &self.output_label { + Some(output_label) => DType::Extension(output_label.clone()), + None => self.storage_dtype.clone(), + } + } + + /// Return how this concrete dispatch executes nullable rows. + pub(crate) fn policy(&self) -> RowPolicy { + self.policy + } + /// Return the output dtype widened with strict input nullability. pub(crate) fn result_dtype(&self, args: &[DType]) -> DType { - let nullability = self.output_dtype.nullability() - | Nullability::from(args.iter().any(DType::is_nullable)); + let output_dtype = self.output_dtype(); + let nullability = + output_dtype.nullability() | Nullability::from(args.iter().any(DType::is_nullable)); + + output_dtype.with_nullability(nullability) + } + + /// Label `values` with the output dtype, preserving their nullability and every value. + /// + /// This is not a cast. [`new`](Self::new) accepted only a label that reinterprets the storage + /// column, so this wraps rather than converts, and an extension dtype takes its nullability + /// from the storage column it wraps. + pub(crate) fn relabel_output(&self, values: ArrayRef) -> VortexResult { + let Some(output_label) = &self.output_label else { + return Ok(values); + }; + + let output_label = output_label.with_nullability(values.dtype().nullability()); + + Ok(ExtensionArray::try_new(output_label, values)?.into_array()) + } + + /// Ensure an executing dispatch reproduced the planned output and policy. + pub(crate) fn ensure_reproduced_by(&self, actual: &Self) -> VortexResult<()> { + vortex_ensure_eq!( + actual.policy, + self.policy, + "row dispatch must select the planned nullable execution policy: planned {:?}, got {:?}", + self.policy, + actual.policy, + ); + vortex_ensure_eq!( + actual.storage_dtype, + self.storage_dtype, + "row dispatch must select the planned storage dtype: planned {}, got {}", + self.storage_dtype, + actual.storage_dtype, + ); + vortex_ensure!( + actual.output_label == self.output_label, + "row dispatch must declare the planned output dtype: planned {}, got {}", + self.output_dtype(), + actual.output_dtype(), + ); + + Ok(()) + } +} + +/// Validate a declared output dtype and return the label to apply to the storage column. +/// +/// A declared dtype **must** be non-nullable and **must** apply by wrapping the finished column, +/// which restricts it to `storage_dtype` itself or to an extension dtype storing exactly that +/// dtype. An extension array holds its storage column as a child, so wrapping costs nothing and +/// works for any encoding. Returning `None` for the former keeps [`BatchPlan::relabel_output`] a +/// no-op for it. +/// +/// This is the single validation point for that relationship, which +/// [`BatchPlan::relabel_output`] relies on. It compares dtypes only. An extension type that +/// constrains its storage values trusts the row kernel to produce values that satisfy it. +fn validate_output_label( + storage_dtype: &DType, + output_dtype: DType, +) -> VortexResult> { + vortex_ensure!( + !output_dtype.is_nullable(), + "a declared row output dtype must be non-nullable, got {output_dtype}", + ); - self.output_dtype.with_nullability(nullability) + if output_dtype == *storage_dtype { + return Ok(None); } + + let DType::Extension(output_label) = output_dtype else { + vortex_bail!( + "a declared row output dtype must be an extension dtype over the storage dtype \ + {storage_dtype}, got {output_dtype}", + ); + }; + vortex_ensure_eq!( + *output_label.storage_dtype(), + *storage_dtype, + "a declared row extension output dtype must store {storage_dtype}, got {}", + output_label.storage_dtype(), + ); + + Ok(Some(output_label)) } /// The nullable execution policy derived from one concrete dispatch. diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs index e7925dcbe62..b4e77a6c50f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs @@ -7,8 +7,11 @@ //! dense attempt containing either unvalidated values or an error that batch execution resolves //! against input validity. +use std::marker::PhantomData; + use vortex_error::VortexResult; +use super::BatchPlan; use super::RowPolicy; use super::RowVisitor; use super::check::assert_deferred_visit_contract; @@ -16,9 +19,9 @@ use super::check::assert_owned_visit_contract; use super::check::assert_sink_visit_contract; use super::check::validate_owned_visit; use super::check::validate_sink_visit; -use super::ensure_plan; use super::row_visitor::private; use crate::ExecutionCtx; +use crate::dtype::DType; use crate::scalar_fn::unstable::row::ElementTuple; use crate::scalar_fn::unstable::row::FailureEvidence; use crate::scalar_fn::unstable::row::IndexedElementTuple; @@ -42,28 +45,40 @@ pub(in crate::scalar_fn::unstable::row) struct ExecuteDenseWithRetry< /// The inputs and planning metadata for this dense attempt. args: &'visit BorrowedRowFnArgs<'inputs>, - /// The function options used to derive a sink's runtime dtype. - options: &'visit F::Options, + /// The output dtype declared by [`RowVisitor::with_output_dtype`], if any. + output_dtype: Option, /// The execution context used to decode the input columns. ctx: &'ctx mut ExecutionCtx, + + /// Ties this visit to the function used by its compile-time contract checks. + function: PhantomData, } impl<'visit, 'inputs, 'ctx, F: RowFn> ExecuteDenseWithRetry<'visit, 'inputs, 'ctx, F> { pub(in crate::scalar_fn::unstable::row) fn new( args: &'visit BorrowedRowFnArgs<'inputs>, - options: &'visit F::Options, ctx: &'ctx mut ExecutionCtx, ) -> Self { - Self { args, options, ctx } + Self { + args, + output_dtype: None, + ctx, + function: PhantomData, + } } } impl private::Sealed for ExecuteDenseWithRetry<'_, '_, '_, F> {} -impl RowVisitor for ExecuteDenseWithRetry<'_, '_, '_, F> { +impl RowVisitor for ExecuteDenseWithRetry<'_, '_, '_, F> { type VisitResult = DenseAttempt; + fn with_output_dtype(mut self, dtype: DType) -> Self { + self.output_dtype = Some(dtype); + self + } + fn visit_prepared( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, @@ -74,12 +89,12 @@ impl RowVisitor for ExecuteDenseWithRetry<'_, '_, '_, F> { Out: OutputElement, { const { assert_owned_visit_contract::() }; - ensure_plan( - self.args.output_dtype(), - self.args.policy(), + let visited = BatchPlan::new( validate_owned_visit::(self.args.dtypes())?, + self.output_dtype, RowPolicy::for_owned_output::(), )?; + self.args.plan().ensure_reproduced_by(&visited)?; execute_owned_infallible::(self.args, self.ctx, prepare, apply) .map(DenseAttempt::Values) @@ -88,29 +103,23 @@ impl RowVisitor for ExecuteDenseWithRetry<'_, '_, '_, F> { fn visit_prepared_into( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - apply: impl Fn( - &Prepared, - Args::Elems<'_>, - >::Row<'_>, - ) -> ApplyResult, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult>::WriteToken>, + Sink: OutputSink, + ApplyResult: SinkResult, { const { assert_sink_visit_contract::() }; - ensure_plan( - self.args.output_dtype(), - self.args.policy(), - validate_sink_visit::(self.options, self.args.dtypes())?, + let visited = BatchPlan::new( + validate_sink_visit::(self.args.dtypes())?, + self.output_dtype, RowPolicy::for_sink::(), )?; + self.args.plan().ensure_reproduced_by(&visited)?; - execute_sink::( - self.args, self.ctx, prepare, apply, - ) - .map(DenseAttempt::Values) + execute_sink::(self.args, self.ctx, prepare, apply) + .map(DenseAttempt::Values) } fn visit_prepared_deferred( @@ -125,12 +134,12 @@ impl RowVisitor for ExecuteDenseWithRetry<'_, '_, '_, F> { Fail: FailureEvidence, { const { assert_deferred_visit_contract::() }; - ensure_plan( - self.args.output_dtype(), - self.args.policy(), + let visited = BatchPlan::new( validate_owned_visit::(self.args.dtypes())?, + self.output_dtype, RowPolicy::for_deferred_output::(), )?; + self.args.plan().ensure_reproduced_by(&visited)?; execute_owned_dense_attempt::( self.args, diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs index 23a0d602d79..4de9e34b330 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -10,6 +10,7 @@ use vortex_error::VortexResult; +use crate::dtype::DType; use crate::scalar_fn::unstable::row::ElementTuple; use crate::scalar_fn::unstable::row::FailureEvidence; use crate::scalar_fn::unstable::row::IndexedElementTuple; @@ -23,14 +24,58 @@ use crate::scalar_fn::unstable::row::SinkResult; /// from constant arguments before visiting any rows. Every visit verifies that the argument tuple /// matches [`RowFn::ARG_NAMES`] and that fallible decoding agrees with [`RowFn::INFALLIBLE`]. /// +/// A visit selects the _storage dtype_, the dtype the chosen [`OutputElement`] or [`OutputSink`] +/// physically builds. [`with_output_dtype`](Self::with_output_dtype) declares the _output dtype_, +/// the dtype the function returns, which defaults to the storage dtype. +/// /// [`RowFn::ARG_NAMES`]: crate::scalar_fn::unstable::row::RowFn::ARG_NAMES /// [`RowFn::INFALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::INFALLIBLE -pub trait RowVisitor: private::Sealed + Sized { +pub trait RowVisitor: private::Sealed + Sized { /// The framework result of visiting one concrete row signature. /// /// This is a batch plan or execution result, not a per-row output. type VisitResult; + /// Declare the dtype this dispatch labels onto the column it builds, replacing the storage + /// dtype as the function's output dtype. + /// + /// Use this for an output dtype derived from the function options or the argument dtypes. A + /// dtype that is a property of the Rust type belongs on [`OutputElement::element_dtype`] or + /// [`OutputSink::storage_dtype`] instead. + /// + /// Batch execution applies the label after deriving nullability and masking null rows, so an + /// empty or all-null batch carries the same metadata as a populated one. + /// + /// `dtype` **must** be non-nullable, and **must** apply by wrapping the finished column. That + /// restricts it to the storage dtype itself or to an extension dtype storing it, because an + /// extension array holds its storage column as a child whatever that column's encoding. The + /// visit validates both, and execution checks that this dispatch declares the dtype planning + /// selected. + /// + /// Labelling validates dtypes and not values. An extension type can constrain its storage + /// values through [`ExtVTable::validate_scalar_value`], so a dispatch that declares one + /// **must** produce values that satisfy it. This is the same trust the framework places in + /// [`OutputElement::element_dtype`]. + /// + /// # Examples + /// + /// Truncate each timestamp to a granularity while preserving its unit and timezone. Deriving + /// the unit once feeds both the output dtype and the row kernel, so the two cannot disagree. + /// + /// ```ignore + /// let (ext_dtype, unit) = timestamp_dtype(&args[0])?; + /// let ticks = options.ticks_in(unit)?; + /// + /// visitor + /// .with_output_dtype(DType::Extension( + /// ext_dtype.with_nullability(Nullability::NonNullable), + /// )) + /// .visit::<(TimestampRow,), i64>(move |(value,)| value - value.rem_euclid(ticks)) + /// ``` + /// [`ExtVTable::validate_scalar_value`]: crate::dtype::extension::ExtVTable::validate_scalar_value + /// [`OutputElement::element_dtype`]: crate::scalar_fn::unstable::row::OutputElement::element_dtype + fn with_output_dtype(self, dtype: DType) -> Self; + /// Visit an infallible row computation that returns one output value per row. /// /// `apply` must not panic or have side effects. Dense execution can pass unspecified values @@ -49,10 +94,10 @@ pub trait RowVisitor: private::Sealed + Sized { /// Dispatch an equality helper over its primitive element type. /// /// ```ignore - /// fn visit_equal(visitor: V) -> VortexResult + /// fn visit_equal(visitor: V) -> VortexResult /// where /// T: NativePType, - /// V: RowVisitor, + /// V: RowVisitor, /// { /// visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_eq(rhs)) /// } @@ -137,12 +182,12 @@ pub trait RowVisitor: private::Sealed + Sized { /// ``` fn visit_into( self, - apply: impl Fn(Args::Elems<'_>, >::Row<'_>) -> ApplyResult, + apply: impl Fn(Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult>::WriteToken>, + Sink: OutputSink, + ApplyResult: SinkResult, { self.visit_prepared_into::( |_| (), @@ -180,16 +225,12 @@ pub trait RowVisitor: private::Sealed + Sized { fn visit_prepared_into( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - apply: impl Fn( - &Prepared, - Args::Elems<'_>, - >::Row<'_>, - ) -> ApplyResult, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult where Args: ElementTuple, - Sink: OutputSink, - ApplyResult: SinkResult>::WriteToken>; + Sink: OutputSink, + ApplyResult: SinkResult; /// Visit a row computation that returns an owned output value and deferred failure evidence. /// diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index e3bdecd0135..2b483d6109b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -96,7 +96,7 @@ pub fn row_fn_return_dtype( ) -> VortexResult { ensure_arity(function, args.len())?; - let plan = function.dispatch(options, args, BatchPlanner::::new(args, options))?; + let plan = function.dispatch(options, args, BatchPlanner::::new(args))?; Ok(plan.result_dtype(args)) } @@ -137,11 +137,13 @@ fn execute_nullary_rows( row_count: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let plan = function.dispatch(options, &[], BatchPlanner::::new(&[], options))?; + let plan = function.dispatch(options, &[], BatchPlanner::::new(&[]))?; let result_dtype = plan.result_dtype(&[]); - let args = BorrowedRowFnArgs::new(&[], row_count, &[], &plan.output_dtype, plan.policy); + let args = BorrowedRowFnArgs::new(&[], row_count, &[], &plan); - let values = execute_row_kernel(function, options, args, ctx)?; + // A nullary function has no input validity to propagate, so its kernel output is the finished + // column and this path labels it directly. + let values = plan.relabel_output(execute_row_kernel(function, options, args, ctx)?)?; finalize_kernel_output(RowFn::id(function), &result_dtype, row_count, values, ctx) } @@ -167,14 +169,7 @@ fn execute_row_kernel( function.dispatch( options, args.dtypes(), - ExecuteRows::::new( - &args, - args.dtypes(), - options, - args.output_dtype(), - args.policy(), - ctx, - ), + ExecuteRows::::new(&args, args.dtypes(), args.plan(), ctx), ) } @@ -187,7 +182,7 @@ fn execute_dense_attempt( function.dispatch( options, args.dtypes(), - ExecuteDenseWithRetry::::new(&args, options, ctx), + ExecuteDenseWithRetry::::new(&args, ctx), ) } @@ -201,15 +196,7 @@ fn try_execute_valid_rows( function.dispatch( options, args.dtypes(), - ExecuteValidRows::::new( - &args, - args.dtypes(), - options, - args.output_dtype(), - args.policy(), - valid, - ctx, - ), + ExecuteValidRows::::new(&args, args.dtypes(), args.plan(), valid, ctx), ) } @@ -219,11 +206,7 @@ fn prepare_batch( args: &dyn ExecutionArgs, ) -> VortexResult { RowFnExecutionArgs::new(RowFn::id(function), args, |arg_dtypes| { - function.dispatch( - options, - arg_dtypes, - BatchPlanner::::new(arg_dtypes, options), - ) + function.dispatch(options, arg_dtypes, BatchPlanner::::new(arg_dtypes)) }) } @@ -282,7 +265,7 @@ mod tests { *ID } - fn dispatch>( + fn dispatch( &self, _options: &Self::Options, _args: &[DType], @@ -304,7 +287,7 @@ mod tests { *ID } - fn dispatch>( + fn dispatch( &self, _options: &Self::Options, args: &[DType], @@ -327,7 +310,7 @@ mod tests { *ID } - fn dispatch>( + fn dispatch( &self, _options: &Self::Options, _args: &[DType],