Let a row dispatch declare a runtime output dtype - #9581
Conversation
`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 <connor.tsui20@gmail.com>
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 <connor.tsui20@gmail.com>
Merging this PR will improve performance by 16.03%
|
| Mode | Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|---|
| ⚡ | Simulation | take[small_m/shuffled/primitive/nonnull/chunks=2048/indices=64] |
801.4 µs | 684 µs | +17.16% |
| ⚡ | Simulation | take[small_m/shuffled/primitive/nonnull/chunks=256/indices=64] |
608.1 µs | 529.3 µs | +14.91% |
Tip
Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.
Comparing claude/generic-output-sink-etzm2l (f3ca34d) with develop (e4b3421)
Footnotes
-
54 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩
| /// The extension dtype labelled onto the finished column, when the declared output dtype | ||
| /// differs from the storage dtype. | ||
| output_label: Option<ExtDTypeRef>, |
There was a problem hiding this comment.
only supports extension types for now, in the future this can be a normal DType that we support more conversions for.
It's kind of similar to our old coercion logic, but a bit more constrained
1 participant
Summary
RowFnAPI #9129Right now, a row function cannot build an output whose dtype depends on its inputs. For example, if a function can accept multiple extension types and returns the same extension type (fn T -> T), the output type needs to know about the input. Right now, the output sink has to do that as well as record how to physically store this data (which may be different from the logical data type).
This PR separates logical and physical types WRT function return. The storage dtype (physical) is what the dispatched
OutputElementorOutputSinkbuilds, and stays a property of the Rust type. The output dtype (logical) is what the function returns, and each dispatch derives it from the options and the argument dtypes.Changes
Adds
RowVisitor::with_output_dtype. A dispatch calls it to declare the dtype that the framework labels (basically just reattaches at the end) onto the column the dispatch builds.Batch execution applies the label in
finalize_output, after it derives nullability and masks the null rows.Sinks become physical only.
OutputSinkloses itsOptionstype parameter, andreturn_dtypebecomes a staticstorage_dtypeinstead.RowVisitoralso losesOptionsparameter.Limitations
There are still limitations. We still do not have a runtime-deciding sink parameterization. For example, if we want a
FixedSizeListSink, we need a way for the sink to be able to represent different row widths. This is fixed in #9585The other is that for types that constrain the values of the storage type (think UUIDv7 or even certain decimal types), we do not validate the storage values are correct (which maybe is fine). Additionally, there will need to be extra logic for things that have more complex wrapping logic (decimal would be the main example), but it's not a huge blocker.
API Changes
OutputSink<Options>becomesOutputSinkreturn_dtype(&Options) -> VortexResult<DType>becomesstorage_dtype() -> DTypeRowVisitor<Options>becomesRowVisitorand gainswith_output_dtype.