From 0ab50e3ebc3fc7cddf75bf363c2aae0d9bb3b613 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 24 Aug 2026 14:17:12 +0000 Subject: [PATCH 1/2] Let a row dispatch declare a runtime output dtype `OutputSink::return_dtype` can read the function options but not the argument dtypes, and `with_capacity` can read neither, so a row function cannot build an output whose exact dtype depends on its inputs. Adds `RowVisitor::with_output_dtype`, which a dispatch calls to declare the dtype it labels onto the column it builds. Planning validates that the label leaves every value unchanged, and batch execution applies it after deriving nullability and masking, so empty and all-null batches carry the same metadata as populated ones. The label reaches every visit method, including the deferred ones that no sink-side design could serve. Sinks become purely physical as a result: `OutputSink` loses its `Options` type parameter, and `return_dtype` becomes a static `storage_dtype`. `RowVisitor` loses the same parameter, which existed only to spell the sink bound. Signed-off-by: Connor Tsui --- .../src/scalar_fn/fns/binary/numeric/row.rs | 6 +- .../src/scalar_fn/unstable/row/batch/args.rs | 36 +- .../unstable/row/batch/execute/mod.rs | 2 +- .../unstable/row/batch/execute/output.rs | 18 +- .../src/scalar_fn/unstable/row/batch/mod.rs | 9 +- .../scalar_fn/unstable/row/batch/planning.rs | 11 +- .../src/scalar_fn/unstable/row/batch/tests.rs | 350 +++++++++++++++++- .../scalar_fn/unstable/row/execute/sink.rs | 74 ++-- .../src/scalar_fn/unstable/row/row_fn.rs | 2 +- .../src/scalar_fn/unstable/row/types/sink.rs | 30 +- .../scalar_fn/unstable/row/visitor/check.rs | 9 +- .../scalar_fn/unstable/row/visitor/execute.rs | 134 ++++--- .../src/scalar_fn/unstable/row/visitor/mod.rs | 25 -- .../scalar_fn/unstable/row/visitor/plan.rs | 208 +++++++++-- .../scalar_fn/unstable/row/visitor/retry.rs | 63 ++-- .../unstable/row/visitor/row_visitor.rs | 59 ++- .../src/scalar_fn/unstable/row/vtable.rs | 43 +-- 17 files changed, 775 insertions(+), 304 deletions(-) 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..d08302ab5be 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 label the storage 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..876687dbe8e 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,160 @@ 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** leave every value unchanged, which +/// restricts it to `storage_dtype` itself or to an extension dtype storing exactly that dtype. +/// 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. +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 label the storage dtype {storage_dtype} without \ + changing any value, 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..bf567e632a4 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,50 @@ 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** label the storage dtype without changing any + /// value. That restricts it to the storage dtype itself or to an extension dtype storing it. + /// The visit validates both, and execution checks that this dispatch declares the dtype + /// planning selected. + /// + /// # 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)) + /// ``` + 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 +86,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 +174,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 +217,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], From f3ca34d3fcd60fe4a106f57bc5931ddeedce9c7f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 24 Aug 2026 15:25:27 +0000 Subject: [PATCH 2/2] Correct why an output label must be an extension dtype The previous wording derived the restriction from value preservation, which does not hold. An extension type can constrain its storage values through `ExtVTable::validate_scalar_value`, and `DivisibleInt` does. The restriction is structural. An extension array holds its storage column as a child, so a label applies by wrapping a finished column of any encoding. No other dtype has that form, so labelling to one would be a cast rather than a wrap. Also records that labelling compares dtypes and does not validate values, which is the same trust the framework already places in `OutputElement::element_dtype`. Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/tests.rs | 2 +- .../src/scalar_fn/unstable/row/visitor/plan.rs | 15 +++++++++------ .../unstable/row/visitor/row_visitor.rs | 16 ++++++++++++---- 3 files changed, 22 insertions(+), 11 deletions(-) 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 d08302ab5be..224c875c52b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -549,7 +549,7 @@ fn test_declared_output_dtype_reaches_planning() -> VortexResult<()> { #[rstest] #[case::nullable(timestamp_dtype(Nullability::Nullable), "must be non-nullable")] -#[case::not_an_extension(DType::from(u64::PTYPE), "must label the storage dtype")] +#[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, 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 876687dbe8e..f603ca6e685 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -236,12 +236,15 @@ impl BatchPlan { /// Validate a declared output dtype and return the label to apply to the storage column. /// -/// A declared dtype **must** be non-nullable and **must** leave every value unchanged, which -/// restricts it to `storage_dtype` itself or to an extension dtype storing exactly that dtype. -/// Returning `None` for the former keeps [`BatchPlan::relabel_output`] a no-op for it. +/// 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. +/// [`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, @@ -257,8 +260,8 @@ fn validate_output_label( let DType::Extension(output_label) = output_dtype else { vortex_bail!( - "a declared row output dtype must label the storage dtype {storage_dtype} without \ - changing any value, got {output_dtype}", + "a declared row output dtype must be an extension dtype over the storage dtype \ + {storage_dtype}, got {output_dtype}", ); }; vortex_ensure_eq!( 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 bf567e632a4..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 @@ -46,10 +46,16 @@ pub trait RowVisitor: private::Sealed + Sized { /// 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** label the storage dtype without changing any - /// value. That restricts it to the storage dtype itself or to an extension dtype storing it. - /// The visit validates both, and execution checks that this dispatch declares the dtype - /// planning selected. + /// `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 /// @@ -66,6 +72,8 @@ pub trait RowVisitor: private::Sealed + Sized { /// )) /// .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.